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.

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

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():
+        try:
-            value = row['value']
+            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 \
+            data_hold['timestamp'] = data['timestamp'].max() if not data.empty else \
-            datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                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_df = DataFrame(data_hold, index=[0])
-        data_hold_melted = data_hold_df.melt(
+            data_hold_melted = data_hold_df.melt(
-            id_vars='timestamp', var_name='variable', value_name='value')
+                id_vars='timestamp', var_name='variable', value_name='value')
-        data_hold_melted['model_id'] = input_data['model_id']
+            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()}",

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()
+    gates_fixture.send_notification.assert_called_once()
-    call_args = gates_fixture.notification_handler.build_and_send_notification.call_args[1]
+    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",

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
+        )

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

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
+        }