diff --git a/scouter/activities/gates.py b/scouter/activities/gates.py index fbdb5e0..cef6bd0 100644 --- a/scouter/activities/gates.py +++ b/scouter/activities/gates.py @@ -16,7 +16,8 @@ quality_gate_filters = { class Gates(BaseActivity): - def apply_aggregation(self, group: DataFrame, aggr_function: str) -> float | None | str: + def apply_aggregation(self, group: DataFrame, aggr_function: str, + metadata: dict[str, Any]) -> float | None | str: """ Apply aggregation function to a group of data. @@ -48,7 +49,8 @@ class Gates(BaseActivity): elif aggr_function == 'min': return group['value'].min() else: - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id="AGGREGATION_ISSUES", message=f"Invalid aggregation function: {aggr_function}", block="aggregate_data", @@ -100,7 +102,8 @@ class Gates(BaseActivity): # Get the latest timestamp latest_timestamp = group['timestamp'].max() - aggr_value = self.apply_aggregation(group, aggr_function) + aggr_value = self.apply_aggregation( + group, aggr_function, metadata) if aggr_value == 'continue': continue @@ -145,7 +148,8 @@ class Gates(BaseActivity): except Exception as e: trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id="AGGREGATION_ISSUES", message=f"Error aggregating data: {e}", block="aggregate_data", @@ -202,7 +206,8 @@ class Gates(BaseActivity): except Exception as e: trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id="DATA_QUALITY_GATE_ISSUES", message=f"Error applying filter {filter_name}: {e}", block="data_quality_gate", @@ -219,7 +224,8 @@ class Gates(BaseActivity): message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}" attachment = filtered_data.to_string() - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}", message=message, block="data_quality_gate", diff --git a/scouter/activities/redis.py b/scouter/activities/redis.py index d9f54e4..0336cb1 100644 --- a/scouter/activities/redis.py +++ b/scouter/activities/redis.py @@ -1,8 +1,10 @@ +import traceback from temporalio import workflow, activity with workflow.unsafe.imports_passed_through(): from logging import Logger from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.activities.redis_base import Redis as RedisBase from sientia_do.temporal.utils.logger import Logger from typing import Any @@ -26,7 +28,18 @@ class Redis(RedisBase): metadata = input_data['metadata'] key = f"last_data_timestamp_{input_data['workflow_name']}_{input_data['schedule_name']}" - data_hold = self.get(key) + try: + data_hold = self.get(key) + except Exception as e: + self.send_notification( + metadata=metadata, + notification_id="REDIS_GET_ERROR", + message=f"Error getting last data timestamp: {e}", + block="get_last_data_timestamp", + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc() + ) + raise e self.debug( f"Last collected timestamp: {data_hold}", @@ -48,6 +61,12 @@ class Redis(RedisBase): data = DataFrame(input_data['data']) + if data.empty: + self.warning("No data to insert", + metadata=metadata + ) + return None + last_data_timestamp = data['inserted_at'].max() self.debug( @@ -55,7 +74,19 @@ class Redis(RedisBase): metadata=metadata ) - self.set(key, last_data_timestamp, ttl=None) + try: + self.set(key, last_data_timestamp, ttl=None) + except Exception as e: + self.send_notification( + metadata=metadata, + notification_id="REDIS_SET_ERROR", + message=f"Error setting last data timestamp: {e}", + + block="put_last_data_timestamp", + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc() + ) + raise e return last_data_timestamp @@ -83,7 +114,18 @@ class Redis(RedisBase): key = f"held_data_{input_data['workflow_name']}_{input_data['schedule_name']}" - data_hold = self.get(key) + try: + data_hold = self.get(key) + except Exception as e: + self.send_notification( + metadata=metadata, + notification_id="REDIS_GET_ERROR", + message=f"Error getting held data: {e}", + block="group_and_hold_data", + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc() + ) + raise e if not data_hold: data_hold = {} @@ -93,22 +135,33 @@ class Redis(RedisBase): ) return data_hold - for _, row in data.iterrows(): - value = row['value'] + try: + for _, row in data.iterrows(): + value = row['value'] - data_hold[row['name']] = value + data_hold[row['name']] = value - data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \ - datetime.now().strftime("%Y-%m-%d %H:%M:%S") + data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \ + datetime.now().strftime("%Y-%m-%d %H:%M:%S") - self.set(key, data_hold, ttl=retention_time) + self.set(key, data_hold, ttl=retention_time) - data_hold_df = DataFrame(data_hold, index=[0]) - data_hold_melted = data_hold_df.melt( - id_vars='timestamp', var_name='variable', value_name='value') - data_hold_melted['model_id'] = input_data['model_id'] + data_hold_df = DataFrame(data_hold, index=[0]) + data_hold_melted = data_hold_df.melt( + id_vars='timestamp', var_name='variable', value_name='value') + data_hold_melted['model_id'] = input_data['model_id'] - data_hold_melted.reset_index(drop=True, inplace=True) + data_hold_melted.reset_index(drop=True, inplace=True) + except Exception as e: + self.send_notification( + metadata=metadata, + notification_id="REDIS_SET_ERROR", + message=f"Error setting held data: {e}", + block="group_and_hold_data", + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc() + ) + raise e self.debug( f"Data grouped and held successfully:\n {data_hold_melted.to_string()}", diff --git a/tests/activities/test_gates.py b/tests/activities/test_gates.py index ba87f5a..366829a 100644 --- a/tests/activities/test_gates.py +++ b/tests/activities/test_gates.py @@ -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", diff --git a/tests/activities/test_mongo.py b/tests/activities/test_mongo.py new file mode 100644 index 0000000..171c6fa --- /dev/null +++ b/tests/activities/test_mongo.py @@ -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 + ) diff --git a/tests/activities/test_redis.py b/tests/activities/test_redis.py index f4d8977..672d035 100644 --- a/tests/activities/test_redis.py +++ b/tests/activities/test_redis.py @@ -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""" diff --git a/tests/utils/test_connectors_config.py b/tests/utils/test_connectors_config.py index bd1694b..e3c4bbd 100644 --- a/tests/utils/test_connectors_config.py +++ b/tests/utils/test_connectors_config.py @@ -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 + }