Merge pull request #8 from Aignosi/SIENTIAPDE-1110-criar-testes-e-2-e

Sientiapde 1110 criar testes e 2 e

requirements.txt

@@ -5,6 +5,6 @@ asyncua
 redis
 aiokafka
 pymongo
-git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.0
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.1
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.1
 pydruid[pandas]

scouter/activities/gates.py

@@ -192,7 +192,8 @@ class Gates(BaseActivity):
             metadata=metadata
         )
 
-        for filter_name, policy in filters.items():
+        for filter_name, config in filters.items():
+            policy = config['policy']
             if filter_name not in quality_gate_filters:
                 self.warning(
                     f"Filter {filter_name} not found",

scouter/activities/redis.py

@@ -93,7 +93,7 @@ class Redis(RedisBase):
     @activity.defn(name="group_and_hold_data")
     async def group_and_hold_data(self, input_data: dict[str, Any]):
         """
-        Groups and holds data in redis. Keep a copy of the most recent 
+        Groups and holds data in redis. Keep a copy of the most recent
         received data for a given pipeline and schedule. This activity updates
         the data in redis and return the full keeped data.
 
@@ -169,3 +169,38 @@ class Redis(RedisBase):
         )
 
         return data_hold_melted.to_dict()
+
+    @activity.defn(name="store_data_package")
+    async def store_data_package(self, input_data: dict[str, Any]):
+        """
+        Stores the data package in redis. It's a debug feature and must be toggled on.
+        input_data:
+            metadata: The metadata of the workflow.
+            workflow_name: The name of the workflow.
+            schedule_name: The name of the schedule.
+            held_data: The final scouter output.
+            data: The data used to collect the data.
+        """
+        metadata = input_data['metadata']
+        key = f"data_package_{input_data['workflow_name']}_{input_data['schedule_name']}_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}"
+
+        data = DataFrame(input_data['data'])
+        held_data = DataFrame(input_data['held_data'])
+
+        cache = {
+            'data': data.to_dict(),
+            'held_data': held_data.to_dict()
+        }
+
+        try:
+            self.set(key, cache, ttl=120)
+        except Exception as e:
+            self.send_notification(
+                metadata=metadata,
+                notification_id="REDIS_SET_ERROR",
+                message=f"Error setting data package: {e}",
+                block="store_data_package",
+                level=NotificationLevel.ERROR,
+                attachment_content=traceback.format_exc()
+            )
+            raise e

scouter/worker/worker.py

@@ -75,6 +75,7 @@ async def main():
                 activities.aggregate_data,
                 activities.group_and_hold_data,
                 activities.export_data_to_postgres,
+                activities.store_data_package
             ]
         ),
         Worker(

scouter/workflow/sub_workflows/core_scouter.py

@@ -74,7 +74,7 @@ class CoreScouter:
         if held_data == {}:
             return
 
-        async_export = workflow.execute_activity_method(
+        await workflow.execute_activity_method(
             Activities.export_data_to_postgres,
             {
                 **metadata,
@@ -86,4 +86,16 @@ class CoreScouter:
             start_to_close_timeout=timedelta(seconds=60)
         )
 
-        await async_export
+        if input_data.get('debug_data_package', False):
+            await workflow.execute_activity_method(
+                Activities.store_data_package,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'held_data': held_data,
+                    'workflow_name': input_data['workflow_name'],
+                    'schedule_name': input_data['schedule_name']
+                },
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=60)
+            )

tests/activities/test_gates.py

@@ -32,7 +32,9 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture):
     # Setup test data
     input_data = {
         'filters': {
-            'NULL_VALUES_FILTER': 'DISCARD'
+            'NULL_VALUES_FILTER': {
+                'policy': 'DISCARD'
+            }
         },
         'data': {
             'tag': ['tag1', 'tag2', 'tag3'],
@@ -62,7 +64,9 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture):
     # Setup test data with out of bounds values
     input_data = {
         'filters': {
-            'OUT_OF_BOUNDS_FILTER': 'KEEP'
+            'OUT_OF_BOUNDS_FILTER': {
+                'policy': 'KEEP'
+            }
         },
         'data': {
             'tag': ['tag1', 'tag2', 'tag3'],
@@ -95,8 +99,12 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture):
     # Setup test data
     input_data = {
         'filters': {
-            'NULL_VALUES_FILTER': 'DISCARD',
+            'NULL_VALUES_FILTER': {
-            'OUT_OF_BOUNDS_FILTER': 'DISCARD'
+                'policy': 'DISCARD'
+            },
+            'OUT_OF_BOUNDS_FILTER': {
+                'policy': 'DISCARD'
+            }
         },
         'data': {
             'tag': ['tag1', 'tag2', 'tag3', 'tag4'],
@@ -129,7 +137,9 @@ async def test_data_quality_gate_with_unknown_filter(gates_fixture):
     gates_fixture.warning = MagicMock()
     input_data = {
         'filters': {
-            'UNKNOWN_FILTER': 'DISCARD'
+            'UNKNOWN_FILTER': {
+                'policy': 'DISCARD'
+            }
         },
         'data': {
             'tag': ['tag1'],
@@ -160,7 +170,9 @@ async def test_data_quality_gate_with_filter_error(gates_fixture):
     # Setup test data
     input_data = {
         'filters': {
-            'NULL_VALUES_FILTER': 'DISCARD'
+            'NULL_VALUES_FILTER': {
+                'policy': 'DISCARD'
+            }
         },
         'data': {
             'tag': ['tag1'],
@@ -199,7 +211,9 @@ async def test_data_quality_gate_with_empty_data(gates_fixture):
     # Setup empty input data
     input_data = {
         'filters': {
-            'NULL_VALUES_FILTER': 'DISCARD'
+            'NULL_VALUES_FILTER': {
+                'policy': 'DISCARD'
+            }
         },
         'data': {
             'tag': [],

tests/activities/test_redis.py

@@ -4,6 +4,7 @@ import pytest
 import numpy as np
 from pandas import DataFrame
 from sientia_do.notifications.handlers import NotificationHandler
+from sientia_do.notifications.models import NotificationLevel
 from scouter.activities.redis import Redis
 
 
@@ -280,3 +281,70 @@ async def test_group_and_hold_data_empty_dataframe(redis_activity):
     result = await redis_activity.group_and_hold_data(test_data)
 
     assert result == {}
+
+
+@pytest.mark.asyncio
+async def test_store_data_package(redis_activity):
+    """Test store_data_package"""
+    redis_activity.set = MagicMock()
+
+    test_data = {
+        **metadata,
+        'workflow_name': 'test_workflow',
+        'schedule_name': 'test_schedule',
+        'held_data': DataFrame({
+            'name': ['sensor1', 'sensor2'],
+            'value': [25.5, 30.0],
+            'timestamp': ['2023-01-01 12:00:00'] * 2
+        }).to_dict(),
+        'data': DataFrame({
+            'name': ['sensor1', 'sensor2'],
+            'value': [25.5, 30.0],
+            'timestamp': ['2023-01-01 12:00:00'] * 2
+        }).to_dict()
+    }
+
+    await redis_activity.store_data_package(test_data)
+
+    redis_activity.set.assert_called_once_with(
+        ANY,
+        {
+            'data': test_data['data'],
+            'held_data': test_data['held_data']
+        },
+        ttl=120)
+
+
+@pytest.mark.asyncio
+async def test_store_data_package_error(redis_activity):
+    """Test store_data_package error"""
+    redis_activity.set = MagicMock(side_effect=Exception('test'))
+    redis_activity.send_notification = MagicMock()
+
+    test_data = {
+        **metadata,
+        'workflow_name': 'test_workflow',
+        'schedule_name': 'test_schedule',
+        'held_data': DataFrame({
+            'name': ['sensor1', 'sensor2'],
+            'value': [25.5, 30.0],
+            'timestamp': ['2023-01-01 12:00:00'] * 2
+        }).to_dict(),
+        'data': DataFrame({
+            'name': ['sensor1', 'sensor2'],
+            'value': [25.5, 30.0],
+            'timestamp': ['2023-01-01 12:00:00'] * 2
+        }).to_dict()
+    }
+
+    with pytest.raises(Exception):
+        await redis_activity.store_data_package(test_data)
+
+    redis_activity.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id="REDIS_SET_ERROR",
+        message="Error setting data package: test",
+        block="store_data_package",
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY
+    )

tests/workflow/sub_workflows/test_core_scouter.py

@@ -34,7 +34,8 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter):
             'schema': 'test_schema',
             'table_name': 'test_table',
             'retention_time': 3600,
-            'model_tags': {}
+            'model_tags': {},
+            'debug_data_package': True
         }
     )
 
@@ -98,6 +99,21 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter):
         )
     ])
 
+    mock_workflow.execute_activity_method.assert_has_calls([
+        call(
+            Activities.store_data_package,
+            {
+                **expected_metadata,
+                'workflow_name': 'test_workflow',
+                'schedule_name': 'test_schedule',
+                'held_data': 'held_data',
+                'data': 'test_data'
+            },
+            retry_policy=ANY,
+            start_to_close_timeout=ANY
+        )
+    ])
+
 
 @pytest.mark.asyncio
 @patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock)