diff --git a/requirements.txt b/requirements.txt index 9f28c99..c8054ba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,5 +3,5 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.1.14 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.2.0 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.38.1 diff --git a/scouter/activities/activities.py b/scouter/activities/activities.py index eec589e..8fc7efb 100644 --- a/scouter/activities/activities.py +++ b/scouter/activities/activities.py @@ -3,10 +3,10 @@ from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.activities.postgres import Postgres from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.temporal.utils.logger import Logger from scouter.activities.redis import Redis from scouter.activities.kafka import Kafka from scouter.activities.gates import Gates - from logging import Logger from typing import Any @@ -62,10 +62,6 @@ class Activities(Postgres, Redis, Kafka, Gates): notification_handler=notification_handler ) - @activity.defn(name="prepare_activity") - async def prepare_activity(self, input_data: dict[str, Any]): - await super().prepare_activity(input_data) - def shutdown(self): Postgres.close(self) Kafka.close(self) diff --git a/scouter/activities/faker.py b/scouter/activities/faker.py index 1ccc4ec..49311ea 100644 --- a/scouter/activities/faker.py +++ b/scouter/activities/faker.py @@ -2,12 +2,12 @@ import random from datetime import datetime, timezone from typing import Any import json -from logging import Logger from kafka import KafkaProducer from temporalio import activity from sientia_do.notifications.handlers import NotificationHandler from sientia_do.temporal.activities.base import BaseActivity +from sientia_do.temporal.utils.logger import Logger class Faker(BaseActivity): @@ -42,6 +42,7 @@ class Faker(BaseActivity): Defaults to random.randint(1, len(self.tags)). """ + metadata = input_data['metadata'] topic = input_data.get('topic') num_messages = input_data.get( 'num_messages', random.randint(1, len(self.tags))) # NOSONAR @@ -49,8 +50,10 @@ class Faker(BaseActivity): if not topic: raise ValueError("Topic must be specified in input_data") - self.logger.info( - f"Generating {num_messages} messages for topic {topic}") + self.info( + f"Generating {num_messages} messages for topic {topic}", + metadata=metadata + ) for _ in range(num_messages): # Select random tag and name @@ -77,4 +80,4 @@ class Faker(BaseActivity): # Ensure all messages are sent self.producer.flush() - self.logger.info("Success") + self.info("Success", metadata=metadata) diff --git a/scouter/activities/gates.py b/scouter/activities/gates.py index 60a66c6..3db376b 100644 --- a/scouter/activities/gates.py +++ b/scouter/activities/gates.py @@ -73,12 +73,16 @@ class Gates(BaseActivity): dict[str, Any]: The aggregated data. """ + metadata = input_data['metadata'] + try: # Convert input data to DataFrame df = DataFrame(input_data['data']) - self.logger.debug( - f"Aggregating time series data: {df.to_string()}") + self.debug( + f"Aggregating time series data: {df.to_string()}", + metadata=metadata + ) # Initialize result dictionary result = {} @@ -101,16 +105,26 @@ class Gates(BaseActivity): if aggr_value == 'continue': continue - self.logger.debug( - f"Aggregated data: {aggr_value}") - self.logger.debug( - f"Latest timestamp: {latest_timestamp}") - self.logger.debug( - f"Groups: {group.to_string()}") - self.logger.debug( - f"group name: {name}") - self.logger.debug( - f"group tag: {tag}") + self.debug( + f"Aggregated data: {aggr_value}", + metadata=metadata + ) + self.debug( + f"Latest timestamp: {latest_timestamp}", + metadata=metadata + ) + self.debug( + f"Groups: {group.to_string()}", + metadata=metadata + ) + self.debug( + f"group name: {name}", + metadata=metadata + ) + self.debug( + f"group tag: {tag}", + metadata=metadata + ) # Store the result result[f"{tag}_{name}"] = { @@ -122,7 +136,10 @@ class Gates(BaseActivity): } result_df = DataFrame(list(result.values())) - self.logger.debug(f"Aggregated data:\n {result_df.to_string()}") + self.debug( + f"Aggregated data:\n {result_df.to_string()}", + metadata=metadata + ) return result_df.to_dict() except Exception as e: @@ -136,7 +153,7 @@ class Gates(BaseActivity): attachment_content=trace ) - self.logger.error(trace) + self.error(trace, metadata=metadata) raise e @activity.defn(name="data_quality_gate") @@ -160,16 +177,23 @@ class Gates(BaseActivity): dict[str, Any]: The data validated. """ + metadata = input_data['metadata'] + filters = input_data['filters'] data = DataFrame(input_data['data']) model_tags = input_data['model_tags'] - self.logger.debug( - f"Applying quality gate to data: {data.to_string()}") + self.debug( + f"Applying quality gate to data: {data.to_string()}", + metadata=metadata + ) for filter_name, policy in filters.items(): if filter_name not in quality_gate_filters: - self.logger.warning(f"Filter {filter_name} not found") + self.warning( + f"Filter {filter_name} not found", + metadata=metadata + ) continue try: @@ -186,7 +210,7 @@ class Gates(BaseActivity): attachment_content=trace ) - self.logger.error(trace) + self.error(trace, metadata=metadata) else: if filtered_data.empty: @@ -206,6 +230,9 @@ class Gates(BaseActivity): if policy == "DISCARD": data = data[~data.index.isin(filtered_data.index)] - self.logger.debug("Data quality gate applied") + self.debug( + "Data quality gate applied", + metadata=metadata + ) return data.to_dict() diff --git a/scouter/activities/kafka.py b/scouter/activities/kafka.py index ee56368..089b28b 100644 --- a/scouter/activities/kafka.py +++ b/scouter/activities/kafka.py @@ -4,6 +4,7 @@ with workflow.unsafe.imports_passed_through(): from logging import Logger from sientia_do.notifications.handlers import NotificationHandler from sientia_do.temporal.activities.base import BaseActivity + from sientia_do.temporal.utils.logger import Logger from typing import Any from kafka import KafkaConsumer from pandas import DataFrame @@ -27,7 +28,7 @@ class Kafka(BaseActivity): def close(self): """Closes the connector connection.""" - self.logger.info("Closing Kafka connector...") + self.info("Closing Kafka connector...") self.kafka_connector.close() def __del__(self): @@ -45,7 +46,12 @@ class Kafka(BaseActivity): dict[str, Any]: The data loaded from the topic. """ - self.logger.debug(f"Loading data from topic: {input_data['topic']}") + metadata = input_data['metadata'] + + self.debug( + f"Loading data from topic: {input_data['topic']}", + metadata=metadata + ) topic = input_data["topic"] @@ -58,7 +64,10 @@ class Kafka(BaseActivity): # Poll for messages records = self.kafka_connector.poll(timeout_ms=self.polling_time) - self.logger.debug(f"Polled {len(records)} records from topic: {topic}") + self.debug( + f"Polled {len(records)} records from topic: {topic}", + metadata=metadata + ) # Process the polled records for _topic_partition, msgs in records.items(): @@ -69,10 +78,14 @@ class Kafka(BaseActivity): if not message_values: return {} - self.logger.debug( - f"Loaded {len(message_values)} messages from topic: {topic}") + self.debug( + f"Loaded {len(message_values)} messages from topic: {topic}", + metadata=metadata + ) - self.logger.debug( - f"Loaded data: {message_values}") + self.debug( + f"Loaded data: {message_values}", + metadata=metadata + ) return DataFrame(message_values).to_dict() diff --git a/scouter/activities/redis.py b/scouter/activities/redis.py index 13b4f54..09c3f40 100644 --- a/scouter/activities/redis.py +++ b/scouter/activities/redis.py @@ -4,6 +4,7 @@ with workflow.unsafe.imports_passed_through(): from logging import Logger from sientia_do.notifications.handlers import NotificationHandler from sientia_do.temporal.activities.redis_base import Redis as RedisBase + from sientia_do.temporal.utils.logger import Logger from typing import Any from pandas import DataFrame from datetime import datetime @@ -31,7 +32,11 @@ class Redis(RedisBase): data (dict[str, Any]): The data to group and hold. retention_time (int): The retention time for data in redis in seconds. """ - self.logger.debug("Grouping and holding data...") + metadata = input_data['metadata'] + + self.debug("Grouping and holding data...", + metadata=metadata + ) data = DataFrame(input_data['data']) retention_time = input_data['retention_time'] @@ -42,7 +47,9 @@ class Redis(RedisBase): if not data_hold: data_hold = {} if data.empty: - self.logger.warning("No data to export") + self.warning("No data to export", + metadata=metadata + ) return data_hold for _, row in data.iterrows(): @@ -62,7 +69,9 @@ class Redis(RedisBase): data_hold_melted.reset_index(drop=True, inplace=True) - self.logger.debug( - f"Data grouped and held successfully:\n {data_hold_melted.to_string()}") + self.debug( + f"Data grouped and held successfully:\n {data_hold_melted.to_string()}", + metadata=metadata + ) return data_hold_melted.to_dict() diff --git a/scouter/workflow/scouter.py b/scouter/workflow/scouter.py index f3c98fa..292834b 100644 --- a/scouter/workflow/scouter.py +++ b/scouter/workflow/scouter.py @@ -32,21 +32,19 @@ class Scouter: input_data['workflow_name'] = 'scouter' - await workflow.execute_local_activity_method( - Activities.prepare_activity, - { - 'workflow_name': input_data['workflow_name'], - 'schedule_name': input_data['schedule_name'], + metadata = { + 'metadata': { + 'model_id': input_data['model_id'], 'model_name': input_data['model_name'], - 'model_id': input_data['model_id'] - }, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60) - ) + 'schedule_name': input_data['schedule_name'], + 'workflow_name': input_data['workflow_name'] + } + } data = await workflow.execute_activity_method( Activities.load_from_kafka, { + **metadata, 'topic': input_data['topic'] }, retry_policy=retry_policy, diff --git a/scouter/workflow/sub_workflows/core_scouter.py b/scouter/workflow/sub_workflows/core_scouter.py index c1a7d2c..9628aea 100644 --- a/scouter/workflow/sub_workflows/core_scouter.py +++ b/scouter/workflow/sub_workflows/core_scouter.py @@ -19,6 +19,7 @@ class CoreScouter: Args: input_data (dict[str, Any]): The data to process. Contains: + metadata (dict[str, Any]): The metadata of the workflow. workflow_name (str): The name of the workflow. schedule_name (str): The name of the schedule. model_name (str): The name of the model. @@ -31,9 +32,19 @@ class CoreScouter: retention_time (int): The retention time for data in redis in seconds. """ + metadata = { + 'metadata': { + 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], + 'schedule_name': input_data['schedule_name'], + 'workflow_name': input_data['workflow_name'] + } + } + filtered_data = await workflow.execute_local_activity_method( Activities.data_quality_gate, { + **metadata, 'filters': input_data['filters'], 'data': input_data['data'], 'model_tags': input_data['model_tags'] @@ -45,6 +56,7 @@ class CoreScouter: grouped_data = await workflow.execute_local_activity_method( Activities.aggregate_data, { + **metadata, 'data': filtered_data, 'model_tags': input_data['model_tags'] }, @@ -55,8 +67,9 @@ class CoreScouter: held_data = await workflow.execute_local_activity_method( Activities.group_and_hold_data, { - 'workflow_name': input_data['workflow_name'], + **metadata, 'schedule_name': input_data['schedule_name'], + 'workflow_name': input_data['workflow_name'], 'data': grouped_data, 'model_id': input_data['model_id'], 'retention_time': input_data['retention_time'] @@ -71,6 +84,7 @@ class CoreScouter: async_export = workflow.execute_activity_method( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': held_data diff --git a/tests/activities/test_activities.py b/tests/activities/test_activities.py index a92e4ca..2150dca 100644 --- a/tests/activities/test_activities.py +++ b/tests/activities/test_activities.py @@ -92,12 +92,14 @@ def test___init__(mock_gates_init, mock_kafka_init, mock_redis_init, mock_postgr ) -@mark.asyncio @patch('scouter.activities.activities.Postgres.__init__') @patch('scouter.activities.activities.Redis.__init__') @patch('scouter.activities.activities.Kafka.__init__') -async def test_prepare_activity(_mock_kafka_init, - _mock_redis_init, _mock_postgres_init): +@patch('scouter.activities.activities.Gates.__init__') +@patch('scouter.activities.activities.Postgres.close') +@patch('scouter.activities.activities.Kafka.close') +def test_shutdown(mock_kafka_close, mock_postgres_close, + _mock_gates_init, _mock_redis_init, _mock_kafka_init, _mock_postgres_init): postgres_config = { 'host': 'localhost', 'port': 5432, @@ -132,20 +134,7 @@ async def test_prepare_activity(_mock_kafka_init, notification_handler=notification_handler ) - input_data = { - 'workflow_name': 'test-workflow-name', - 'schedule_name': 'test-schedule-name', - 'model_name': 'test-model-name', - 'model_id': 'test-model-id' - } + activities.shutdown() - await activities.prepare_activity(input_data) - - assert activities.notification_handler.base_notification.pipeline == input_data[ - 'workflow_name'] - assert activities.notification_handler.base_notification.trigger == input_data[ - 'schedule_name'] - assert activities.notification_handler.base_notification.model_name == input_data[ - 'model_name'] - assert activities.notification_handler.base_notification.model_id == input_data[ - 'model_id'] + mock_postgres_close.assert_called_once() + mock_kafka_close.assert_called_once() diff --git a/tests/activities/test_faker.py b/tests/activities/test_faker.py index 5a35abf..40f4893 100644 --- a/tests/activities/test_faker.py +++ b/tests/activities/test_faker.py @@ -1,4 +1,3 @@ -from logging import Logger from unittest.mock import MagicMock, patch, call import pytest from sientia_do.notifications.handlers import NotificationHandler @@ -22,7 +21,7 @@ def mock_datetime(): @pytest.fixture def faker_instance(mock_kafka_producer): - logger = MagicMock(spec=Logger) + logger = MagicMock() notification_handler = MagicMock(spec=NotificationHandler) return Faker( bootstrap_servers='localhost:9092', @@ -37,6 +36,15 @@ async def test_faker_init(faker_instance, mock_kafka_producer): assert faker_instance.producer is not None assert len(faker_instance.tags) == 6 +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } +} + @pytest.mark.asyncio async def test_generate_and_send_data_default_count(faker_instance, @@ -56,7 +64,7 @@ async def test_generate_and_send_data_default_count(faker_instance, ] # Call the method - await faker_instance.generate_and_send_data({'topic': 'test_topic'}) + await faker_instance.generate_and_send_data({'topic': 'test_topic', **metadata}) # Verify the producer was called 3 times (default count) assert mock_kafka_producer.send.call_count == 3 @@ -92,7 +100,8 @@ async def test_generate_and_send_data_custom_count(faker_instance, mock_kafka_pr # Call the method with custom count await faker_instance.generate_and_send_data({ 'topic': 'test_topic', - 'num_messages': 2 + 'num_messages': 2, + **metadata }) # Verify the producer was called 2 times @@ -104,7 +113,7 @@ async def test_generate_and_send_data_custom_count(faker_instance, mock_kafka_pr async def test_generate_and_send_data_no_topic(faker_instance): """Test that ValueError is raised when no topic is provided""" with pytest.raises(ValueError, match="Topic must be specified in input_data"): - await faker_instance.generate_and_send_data({}) + await faker_instance.generate_and_send_data({**metadata}) @pytest.mark.asyncio @@ -112,7 +121,7 @@ async def test_generate_and_send_data_no_topic(faker_instance): async def test_generate_and_send_data_random_values(_random, faker_instance, mock_kafka_producer): """Test that random values are within expected ranges""" # Call the method - await faker_instance.generate_and_send_data({'topic': 'test_topic'}) + await faker_instance.generate_and_send_data({'topic': 'test_topic', **metadata}) # Get the call arguments call_args = mock_kafka_producer.send.call_args[1]['value'] @@ -132,7 +141,8 @@ async def test_generate_and_send_data_generate_null_values( mock_kafka_producer): # Call the method await faker_instance.generate_and_send_data({'topic': 'test_topic', - 'num_messages': 1}) + 'num_messages': 1, + **metadata}) # Get the call arguments call_args = mock_kafka_producer.send.call_args[1]['value'] diff --git a/tests/activities/test_gates.py b/tests/activities/test_gates.py index 9569d57..245b176 100644 --- a/tests/activities/test_gates.py +++ b/tests/activities/test_gates.py @@ -14,6 +14,16 @@ def gates_fixture(): return Gates(logger=logger, notification_handler=notification_handler) +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } +} + + @pytest.mark.asyncio async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture): """Test data_quality_gate with NULL_VALUES_FILTER and DISCARD policy.""" @@ -31,7 +41,8 @@ async def test_data_quality_gate_with_null_values_filter_discard(gates_fixture): 'tag1': {'data_range': [0, 100]}, 'tag2': {'data_range': [0, 100]}, 'tag3': {'data_range': [0, 100]} - } + }, + **metadata } # Execute @@ -60,7 +71,8 @@ async def test_data_quality_gate_with_out_of_bounds_filter_keep(gates_fixture): 'tag1': {'data_range': [0, 100]}, 'tag2': {'data_range': [0, 100]}, 'tag3': {'data_range': [0, 100]} - } + }, + **metadata } # Mock the out_of_bounds_filter to return rows with out of bounds values @@ -95,7 +107,8 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture): 'tag2': {'data_range': [0, 100]}, 'tag3': {'data_range': [0, 100]}, 'tag4': {'data_range': [0, 100]} - } + }, + **metadata } result = await gates_fixture.data_quality_gate(input_data) @@ -111,6 +124,7 @@ async def test_data_quality_gate_with_multiple_filters(gates_fixture): async def test_data_quality_gate_with_unknown_filter(gates_fixture): """Test data_quality_gate with an unknown filter.""" # Setup test data with unknown filter + gates_fixture.warning = MagicMock() input_data = { 'filters': { 'UNKNOWN_FILTER': 'DISCARD' @@ -123,7 +137,8 @@ async def test_data_quality_gate_with_unknown_filter(gates_fixture): }, 'model_tags': { 'tag1': {'data_range': [0, 100]} - } + }, + **metadata } # Execute @@ -131,8 +146,10 @@ async def test_data_quality_gate_with_unknown_filter(gates_fixture): # Verify data is unchanged and warning is logged assert len(result['tag']) == 1 - gates_fixture.logger.warning.assert_called_once_with( - "Filter UNKNOWN_FILTER not found") + gates_fixture.warning.assert_called_once_with( + "Filter UNKNOWN_FILTER not found", + metadata=metadata['metadata'] + ) @pytest.mark.asyncio @@ -151,7 +168,8 @@ async def test_data_quality_gate_with_filter_error(gates_fixture): }, 'model_tags': { 'tag1': {'data_range': [0, 100]} - } + }, + **metadata } # Mock the filter to raise an exception @@ -187,7 +205,8 @@ async def test_data_quality_gate_with_empty_data(gates_fixture): 'value': [], 'timestamp': [] }, - 'model_tags': {} + 'model_tags': {}, + **metadata } # Execute @@ -212,7 +231,8 @@ async def test_data_quality_gate_with_no_filters(gates_fixture): }, 'model_tags': { 'tag1': {'data_range': [0, 100]} - } + }, + **metadata } # Execute @@ -272,7 +292,8 @@ async def test_aggregate_data(gates_fixture): 'model_tags': { 'name1': {'aggr_function': 'avg'}, 'name2': {'aggr_function': 'max'}, - } + }, + **metadata } # Expected result @@ -309,7 +330,8 @@ async def test_aggregate_data_with_continue(gates_fixture): 'model_tags': { 'name1': {'aggr_function': 'avg'}, 'name2': {'aggr_function': 'max'}, - } + }, + **metadata } # Expected result @@ -343,7 +365,8 @@ async def test_aggregate_data_raise_exception(gates_fixture): 'model_tags': { 'name1': {'aggr_function': 'avg'}, 'name2': {'aggr_function': 'max'}, - } + }, + **metadata } try: diff --git a/tests/activities/test_kafka.py b/tests/activities/test_kafka.py index d38abfd..681c9dd 100644 --- a/tests/activities/test_kafka.py +++ b/tests/activities/test_kafka.py @@ -38,9 +38,19 @@ def test___init__(kafka_consumer): ) +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } +} + + @mark.asyncio async def test_load_from_kafka(kafka): - input_data = {"topic": "test-topic"} + input_data = {"topic": "test-topic", **metadata} data = [ ("test-topic", [ @@ -67,7 +77,7 @@ async def test_load_from_kafka(kafka): @mark.asyncio async def test_load_from_kafka_empty(kafka): - input_data = {"topic": "test-topic"} + input_data = {"topic": "test-topic", **metadata} kafka.kafka_connector.poll.return_value = MagicMock( items=MagicMock(return_value=[]) diff --git a/tests/activities/test_redis.py b/tests/activities/test_redis.py index b86a04d..a266c71 100644 --- a/tests/activities/test_redis.py +++ b/tests/activities/test_redis.py @@ -41,11 +41,22 @@ def test_redis_initialization(mock_redis_init): ) +metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } +} + + @pytest.mark.asyncio async def test_group_and_hold_data_new_key(redis_activity): """Test group_and_hold_data with a new key""" # Setup test_data = { + **metadata, 'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule', 'retention_time': 3600, @@ -97,6 +108,7 @@ async def test_group_and_hold_data_update_existing(redis_activity): # New data to update with test_data = { + **metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', 'retention_time': 3600, @@ -142,6 +154,7 @@ async def test_group_and_hold_data_with_none_values(redis_activity): """Test handling of None values in group_and_hold_data""" # Setup test data with None values test_data = { + **metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', 'retention_time': 3600, @@ -170,6 +183,7 @@ async def test_group_and_hold_data_empty_dataframe(redis_activity): """Test group_and_hold_data with empty DataFrame""" # Setup test with empty data test_data = { + **metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', 'retention_time': 3600, diff --git a/tests/workflow/sub_workflows/test_core_scouter.py b/tests/workflow/sub_workflows/test_core_scouter.py index a87d063..905b864 100644 --- a/tests/workflow/sub_workflows/test_core_scouter.py +++ b/tests/workflow/sub_workflows/test_core_scouter.py @@ -30,10 +30,20 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): } ) + expected_metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'test_workflow' + } + } + mock_workflow.execute_local_activity_method.assert_has_calls([ call( Activities.data_quality_gate, { + **expected_metadata, 'filters': {'test_filter': 'test_value'}, 'data': 'test_data', 'model_tags': {} @@ -45,6 +55,7 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): call( Activities.aggregate_data, { + **expected_metadata, 'data': 'filtered_data', 'model_tags': {} }, @@ -55,6 +66,7 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): call( Activities.group_and_hold_data, { + **expected_metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', 'data': 'grouped_data', @@ -69,6 +81,7 @@ async def test_core_scouter_workflow_success(mock_workflow, core_scouter): call( Activities.export_data_to_postgres, { + **expected_metadata, 'schema': 'test_schema', 'table_name': 'test_table', 'data': 'held_data'}, @@ -98,10 +111,20 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter } ) + expected_metadata = { + 'metadata': { + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'schedule_name': 'test_schedule', + 'workflow_name': 'test_workflow' + } + } + mock_workflow.execute_local_activity_method.assert_has_calls([ call( Activities.data_quality_gate, { + **expected_metadata, 'filters': {'test_filter': 'test_value'}, 'data': 'test_data', 'model_tags': {} @@ -113,6 +136,7 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter call( Activities.aggregate_data, { + **expected_metadata, 'data': {}, 'model_tags': {} }, @@ -123,6 +147,7 @@ async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter call( Activities.group_and_hold_data, { + **expected_metadata, 'workflow_name': 'test_workflow', 'schedule_name': 'test_schedule', 'data': {}, diff --git a/tests/workflow/test_scouter.py b/tests/workflow/test_scouter.py index f1402c7..9727293 100644 --- a/tests/workflow/test_scouter.py +++ b/tests/workflow/test_scouter.py @@ -23,21 +23,19 @@ async def test_scouter_workflow(mock_workflow, scouter): } ) - mock_workflow.execute_local_activity_method.assert_called_once_with( - Activities.prepare_activity, - { - 'workflow_name': 'scouter', - 'schedule_name': 'test_schedule', + expected_metadata = { + 'metadata': { + 'model_id': 'test_model_id', 'model_name': 'test_model', - 'model_id': 'test_model_id' - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } + } mock_workflow.execute_activity_method.assert_called_once_with( Activities.load_from_kafka, { + **expected_metadata, 'topic': 'test_topic' }, retry_policy=ANY, @@ -70,21 +68,19 @@ async def test_scouter_workflow_empty(mock_workflow, scouter): } ) - mock_workflow.execute_local_activity_method.assert_called_once_with( - Activities.prepare_activity, - { - 'workflow_name': 'scouter', - 'schedule_name': 'test_schedule', + expected_metadata = { + 'metadata': { + 'model_id': 'test_model_id', 'model_name': 'test_model', - 'model_id': 'test_model_id' - }, - retry_policy=ANY, - start_to_close_timeout=ANY - ) + 'schedule_name': 'test_schedule', + 'workflow_name': 'scouter' + } + } mock_workflow.execute_activity_method.assert_called_once_with( Activities.load_from_kafka, { + **expected_metadata, 'topic': 'test_topic' }, retry_policy=ANY,