SIENTIAPDE-1110

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

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",

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
)

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"""

View File

@@ -2,6 +2,8 @@ import os
from unittest.mock import patch
import pytest
from scouter.utils.connectors_config import (
build_druid_config,
build_mongodb_config,
build_postgres_config,
build_kafka_config,
build_redis_config
@@ -113,3 +115,53 @@ def test_build_redis_config_with_env_vars():
'username': 'test',
'password': 'test'
}
def test_build_mongodb_config_defaults():
"""Test that build_mongodb_config returns default values when no env vars are set"""
config = build_mongodb_config()
assert config == {
'connection_string': 'mongodb://sientia:sientia@localhost:27017', # NOSONAR
'database_name': 'sientia'
}
def test_build_mongodb_config_with_env_vars():
"""Test that build_mongodb_config uses env vars when set"""
with patch.dict(os.environ, {
'MONGODB_URL': 'mongodb.example.com:27017',
'MONGODB_DATABASE_NAME': 'test_db',
'MONGODB_USERNAME': 'test',
'MONGODB_PASSWORD': 'test'
}):
config = build_mongodb_config()
assert config == {
'connection_string': 'mongodb://test:test@mongodb.example.com:27017',
'database_name': 'test_db'
}
def test_build_druid_config_defaults():
"""Test that build_druid_config returns default values when no env vars are set"""
config = build_druid_config()
assert config == {
'host': 'localhost',
'port': 8082
}
def test_build_druid_config_with_env_vars():
"""Test that build_druid_config uses env vars when set"""
with patch.dict(os.environ, {
'DRUID_HOST': 'druid.example.com',
'DRUID_PORT': '8083'
}):
config = build_druid_config()
assert config == {
'host': 'druid.example.com',
'port': 8083
}