gitea_admin/sientia-dataops-scouter_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1110

Browse Source 
Enhance Gates and Redis activities by adding metadata parameter to apply_aggregation and notification methods. Refactor notification handling to use send_notification for improved consistency. Update tests to reflect changes in notification method calls and ensure proper functionality with new metadata integration.
This commit is contained in:
vitor-aignosi
2025-07-04 11:01:20 -03:00
parent d08d1b1337
commit 08240efc8b
6 changed files with 425 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
tests/activities/test_gates.py
View File

@@ -11,7 +11,9 @@ def gates_fixture():
"""Fixture to create a Gates instance with mocked dependencies."""
logger = Mock()
notification_handler = MagicMock()
return Gates(logger=logger, notification_handler=notification_handler)
gates = Gates(logger=logger, notification_handler=notification_handler)
gates.send_notification = MagicMock()
return gates
metadata = {
@@ -51,7 +53,7 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
# Verify
assert len(result['tag']) == 2
assert 'tag2' not in result['tag']
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
@pytest.mark.asyncio
@@ -84,7 +86,7 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture):
# Verify data is kept but notification is sent
assert len(result['tag']) == 3 # All rows kept
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
@pytest.mark.asyncio
@@ -117,7 +119,7 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture):
assert result == {'tag': {0: 'tag1', 3: 'tag4'}, 'name': {0: 'tag1', 3: 'tag4'}, 'value': {
0: 1.0, 3: 4.0}, 'timestamp': {0: '2023-01-01', 3: '2023-01-04'}}
# Should be called twice (once for each filter)
assert gates_fixture.notification_handler.build_and_send_notification.call_count == 2
assert gates_fixture.send_notification.call_count == 2
@pytest.mark.asyncio
@@ -184,8 +186,8 @@ async def test_data_quality_gate_with_filter_error(gates_fixture):
# Verify error notification is sent and data is unchanged
assert len(result['tag']) == 1
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
call_args = gates_fixture.notification_handler.build_and_send_notification.call_args[1]
gates_fixture.send_notification.assert_called_once()
call_args = gates_fixture.send_notification.call_args[1]
assert call_args['notification_id'] == "DATA_QUALITY_GATE_ISSUES"
assert call_args['level'] == NotificationLevel.ERROR
assert "Filter error" in call_args['message']
@@ -214,7 +216,7 @@ async def test_data_quality_gate_with_empty_data(gates_fixture):
# Verify empty result and no notifications
assert len(result['tag']) == 0
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -240,7 +242,7 @@ async def test_data_quality_gate_with_no_filters(gates_fixture):
# Verify data is unchanged and no notifications
assert len(result['tag']) == 1
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.parametrize(
@@ -265,14 +267,15 @@ async def test_data_quality_gate_with_no_filters(gates_fixture):
)
def test_apply_aggregation(gates_fixture, group_data, aggr_function, expected_result):
"""Test apply_aggregation method with various scenarios."""
result = gates_fixture.apply_aggregation(group_data, aggr_function)
result = gates_fixture.apply_aggregation(
group_data, aggr_function, metadata)
assert result == expected_result
# Check notification was sent for invalid function
if aggr_function == 'invalid':
gates_fixture.notification_handler.build_and_send_notification.assert_called_once()
gates_fixture.send_notification.assert_called_once()
else:
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -308,7 +311,7 @@ async def test_aggregate_data(gates_fixture):
# Verify
assert result == expected_result
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -342,7 +345,7 @@ async def test_aggregate_data_with_continue(gates_fixture):
# Verify
assert result == expected_result
gates_fixture.notification_handler.build_and_send_notification.assert_not_called()
gates_fixture.send_notification.assert_not_called()
@pytest.mark.asyncio
@@ -373,7 +376,8 @@ async def test_aggregate_data_raise_exception(gates_fixture):
await gates_fixture.aggregate_data(input_data)
except Exception as e:
assert str(e) == "Test exception"
gates_fixture.notification_handler.build_and_send_notification.assert_called_once_with(
gates_fixture.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id="AGGREGATION_ISSUES",
message="Error aggregating data: Test exception",
block="aggregate_data",

192
tests/activities/test_mongo.py Normal file
View File

@@ -0,0 +1,192 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from scouter.activities.mongodb import MongoDB, clear_mongo_id
def test_clear_mongo_id():
"""Test clear_mongo_id"""
data = [
{'_id': '1', 'name': 'test1'},
{'_id': '2', 'name': [{
'_id': '3',
'name': 'test3'
}]}
]
result = clear_mongo_id(data)
assert result == [{'name': 'test1'}, {'name': [{'name': 'test3'}]}]
@patch('scouter.activities.mongodb.MongoClient')
def test_mongodb___init__(mock_mongo_client):
"""Test MongoDB __init__"""
mongo = MongoDB(
connection_string='mongodb://localhost:27017',
database_name='test_db',
logger=MagicMock(),
notification_handler=MagicMock()
)
mock_mongo_client.assert_called_once_with(
'mongodb://localhost:27017',
serverSelectionTimeoutMS=5000
)
mock_mongo_client.return_value.server_info.assert_called_once()
mock_mongo_client.return_value.__getitem__.assert_called_once_with(
'test_db')
assert mongo.client is not None
assert mongo.database is not None
@fixture
@patch('scouter.activities.mongodb.MongoClient')
def mongodb_activity(mock_mongo_client):
"""Test MongoDB activity"""
mongo = MongoDB(
connection_string='mongodb://localhost:27017',
database_name='test_db',
logger=MagicMock(),
notification_handler=MagicMock()
)
return mongo
def test_shutdown_success(mongodb_activity):
"""Test shutdown"""
mongodb_activity.shutdown()
mongodb_activity.client.close.assert_called_once()
def test_shutdown_error(mongodb_activity):
"""Test shutdown"""
mongodb_activity.client.close = MagicMock(side_effect=Exception('test'))
mongodb_activity.shutdown()
mongodb_activity.client.close.assert_called_once()
@mark.asyncio
async def test_load_latest_data_none_last_data_timestamp(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.return_value = [
{
'name': 'test1',
'value': 1,
'inserted_at': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
]
result = await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': None
})
mongodb_activity.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{},
{"_id": 0}
)
assert result == {
'name': {
0: 'test1'
},
'value': {
0: 1
},
'inserted_at': {
0: '2023-01-01 12:00:00.000000'
}
}
@mark.asyncio
async def test_load_latest_data_not_none_last_data_timestamp(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.return_value = [
{
'name': 'test1',
'value': 1,
'inserted_at': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
]
result = await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000'
})
mongodb_activity.database.__getitem__.assert_called_once_with(
'test_collection')
collection.find.assert_called_once_with(
{
'inserted_at': {
'$gt': datetime.strptime(
'2023-01-01 12:00:00.000000', '%Y-%m-%d %H:%M:%S.%f')
}
},
{"_id": 0}
)
assert result == {
'name': {
0: 'test1'
},
'value': {
0: 1
},
'inserted_at': {
0: '2023-01-01 12:00:00.000000'
}
}
@mark.asyncio
async def test_load_latest_data_error(mongodb_activity):
"""Test load_latest_data"""
collection = MagicMock()
mongodb_activity.send_notification = MagicMock()
mongodb_activity.database.__getitem__.return_value = collection
collection.find.side_effect = Exception('test')
try:
await mongodb_activity.load_latest_data({
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
'collection_name': 'test_collection',
'last_data_timestamp': '2023-01-01 12:00:00.000000'
})
except Exception as e:
assert str(e) == 'test'
mongodb_activity.send_notification.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test',
block='load_latest_data',
level=NotificationLevel.ERROR,
attachment_content=ANY
)

84
tests/activities/test_redis.py
View File

@@ -51,6 +51,90 @@ metadata = {
}
@pytest.mark.asyncio
async def test_get_last_data_timestamp_none(redis_activity):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'
}
redis_activity.get = MagicMock(return_value=None)
result = await redis_activity.get_last_data_timestamp(test_data)
assert result is None
@pytest.mark.asyncio
async def test_get_last_data_timestamp_not_none(redis_activity):
"""Test get_last_data_timestamp"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule'
}
redis_activity.get = MagicMock(return_value='2023-01-01 12:00:00')
result = await redis_activity.get_last_data_timestamp(test_data)
redis_activity.get.assert_called_once_with(
'last_data_timestamp_test_pipeline_test_schedule'
)
assert result == '2023-01-01 12:00:00'
@pytest.mark.asyncio
async def test_put_last_data_timestamp_empty_dataframe(redis_activity):
"""Test put_last_data_timestamp with empty dataframe"""
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule',
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records')
}
redis_activity.set = MagicMock()
result = await redis_activity.put_last_data_timestamp(test_data)
assert result is None
redis_activity.set.assert_not_called()
@pytest.mark.asyncio
async def test_put_last_data_timestamp_not_empty_dataframe(redis_activity):
"""Test put_last_data_timestamp with not empty dataframe"""
data = DataFrame({
'name': ['sensor1', 'sensor2'],
'value': [25.5, 30.0],
'inserted_at': ['2023-01-01 12:00:00', '2023-01-01 12:00:01']
})
test_data = {
**metadata,
'workflow_name': 'test_pipeline',
'schedule_name': 'test_schedule',
'data': data.to_dict('records')
}
redis_activity.set = MagicMock()
result = await redis_activity.put_last_data_timestamp(test_data)
assert result == '2023-01-01 12:00:01'
redis_activity.set.assert_called_once_with(
'last_data_timestamp_test_pipeline_test_schedule',
'2023-01-01 12:00:01',
ttl=None
)
@pytest.mark.asyncio
async def test_group_and_hold_data_new_key(redis_activity):
"""Test group_and_hold_data with a new key"""