diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 39199d9..6147382 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -90,7 +90,8 @@ class Gates(BaseActivity): filter_output.append(config['policy']) except Exception as e: trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id=f"INTPUT_GATE_ERROR__{fil}", message=f"Error in filter {fil}:{config}: \n {e}", block="input_gate", @@ -145,7 +146,8 @@ class Gates(BaseActivity): if mlflow_response_filter_functions[fil](data, config): filter_output.append(config['policy']) comments.append(data['content']['message']) - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}", message=data['content']['message'], block="mlflow_gate", @@ -154,7 +156,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=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}", message=f"Error in filter {fil}:{config}: \n {e}", block="mlflow_gate", @@ -208,7 +211,8 @@ class Gates(BaseActivity): try: if mlflow_content_filter_functions[fil](data, config): filter_output.append(config['policy']) - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}", message=f"Data not passed the content filter {fil}:{config}", block="mlflow_gate", @@ -217,7 +221,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=f"MLFLOW_GATE_CONTENT_FILTER__{fil}", message=f"Error in filter {fil}:{config}: \n {e}", block="mlflow_gate", diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 93fab25..6b62e22 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -60,7 +60,8 @@ class OPC(BaseActivity): tag, data, data_type, self.logger, metadata) except Exception as e: trace = traceback.format_exc() - self.notification_handler.build_and_send_notification( + self.send_notification( + metadata=metadata, notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR", message=f"Error writing data to OPC server: {e}", block="write_opc_data", diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 9208559..f0a768e 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/laborious/utils/repository/opc_repository.py @@ -217,7 +217,7 @@ class OpcRepository(): trace = traceback.format_exc() self.notification_handler.build_and_send_notification( notification_id=f"OPC_WRITE_GET_NODE_ERROR_{self.id}", - message=f"Failed to get node from OPC server: {e}", + message=f"Failed to get node from OPC server: {e} | metadata: {metadata}", block="opc_repository", level=NotificationLevel.ERROR, attachment_content=trace @@ -229,7 +229,7 @@ class OpcRepository(): if data_type not in data_type_map: self.notification_handler.build_and_send_notification( notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}", - message=f"Unsupported data type: {data_type}", + message=f"Unsupported data type: {data_type} | metadata: {metadata}", block="opc_repository", level=NotificationLevel.ERROR ) @@ -258,7 +258,7 @@ class OpcRepository(): trace = traceback.format_exc() self.notification_handler.build_and_send_notification( notification_id=f"OPC_WRITE_DATA_ERROR_{self.id}", - message=f"Failed to write data to OPC server: {e}", + message=f"Failed to write data to OPC server: {e} | metadata: {metadata}", block="opc_repository", level=NotificationLevel.ERROR, attachment_content=trace diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 31f83e8..e482ca8 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -90,65 +90,6 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre ) -@mark.asyncio -@patch('laborious.activities.activities.Postgres.__init__') -@patch('laborious.activities.activities.MLFlow.__init__') -@patch('laborious.activities.activities.OPC.__init__') -async def test_prepare_activity(_mock_opc_init, - _mock_mlflow_init, _mock_postgres_init): - postgres_config = { - 'host': 'localhost', - 'port': 5432, - 'user': 'postgres', - 'password': 'postgres', - 'dbname': 'postgres', - 'min_connections': 1, - 'max_connections': 10 - } - - mlflow_config = { - 'host': 'localhost', - 'port': 5000, - 'username': 'mlflow', - 'password': 'mlflow' - } - - opc_config = { - 'bootstrap_servers': 'localhost:9092', - 'polling_time': 1000, - 'group_id': 'test-group' - } - - logger = MagicMock() - notification_handler = MagicMock() - - activities = Activities( - postgres_config=postgres_config, - mlflow_config=mlflow_config, - opc_config=opc_config, - logger=logger, - 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' - } - - 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'] - - @patch('laborious.activities.activities.Postgres', return_value=MagicMock()) @patch('laborious.activities.activities.MLFlow', return_value=MagicMock()) @patch('laborious.activities.activities.OPC', return_value=MagicMock()) diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 28ec1d2..19d52cd 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -6,16 +6,34 @@ from laborious.activities.gates import Gates @fixture def gates_activity(): - return Gates( + gates = Gates( logger=MagicMock(), notification_handler=MagicMock(), ) + gates.error = MagicMock() + gates.debug = MagicMock() + gates.info = MagicMock() + gates.warning = MagicMock() + gates.critical = MagicMock() + gates.send_notification = MagicMock() + return gates + + +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} @mark.asyncio async def test_input_gate_invalid_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'INVALID_FILTER': {'POLICY': 'STOP'} }, @@ -28,8 +46,8 @@ async def test_input_gate_invalid_filter(gates_activity): # Assert assert result == (None, 0, "") - gates_activity.logger.error.assert_called_once_with( - "Filter INVALID_FILTER not found" + gates_activity.error.assert_called_once_with( + "Filter INVALID_FILTER not found", metadata['metadata'] ) @@ -41,6 +59,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac mock_input_filter_functions.__getitem__.return_value = MagicMock( side_effect=Exception("Test error")) input_data = { + **metadata, 'filters': { 'EMPTY_DATA': {'policy': 'STOP', 'config': {}} }, @@ -53,7 +72,8 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac # Assert assert result == (None, 0, "") - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( + gates_activity.send_notification.assert_called_once_with( + metadata=metadata['metadata'], notification_id="INTPUT_GATE_ERROR__EMPTY_DATA", message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error", block="input_gate", @@ -66,6 +86,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac async def test_input_gate_no_filters(gates_activity): # Arrange input_data = { + **metadata, 'filters': {}, 'data': {'value': [1, 2, 3]}, 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'] @@ -76,13 +97,14 @@ async def test_input_gate_no_filters(gates_activity): # Assert assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_input_gate_with_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'EMPTY_DATA': {'policy': 'STOP', 'config': {}} }, @@ -95,13 +117,14 @@ async def test_input_gate_with_filter(gates_activity): # Assert assert result == ('STOP', -1, "Input data with bad quality") - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_mlflow_response_gate_invalid_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'INVALID_FILTER': {'POLICY': 'STOP'} }, @@ -126,6 +149,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock( side_effect=Exception("Test error")) input_data = { + **metadata, 'filters': { 'INVALID_FILTER': {'POLICY': 'STOP'} }, @@ -139,7 +163,8 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter # Assert assert result == (None, 0, "") - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( + gates_activity.send_notification.assert_called_once_with( + metadata=metadata['metadata'], notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER", message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error", block="mlflow_gate", @@ -152,6 +177,7 @@ async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter async def test_mlflow_response_gate_no_filters(gates_activity): # Arrange input_data = { + **metadata, 'filters': {}, 'data': {'content': {'message': 'success'}}, 'type': 'test', @@ -163,13 +189,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity): # Assert assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_mlflow_response_gate_with_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'API_ERROR': {'policy': 'STOP'} }, @@ -189,14 +216,15 @@ async def test_mlflow_response_gate_with_filter(gates_activity): # Assert assert result == ('STOP', -1, "API error occurred") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called() + gates_activity.debug.assert_called() + gates_activity.send_notification.assert_called() @mark.asyncio async def test_mlflow_content_gate_invalid_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'INVALID_FILTER': {'POLICY': 'STOP'} }, @@ -221,6 +249,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock( side_effect=Exception("Test error")) input_data = { + **metadata, 'filters': { 'API_ERROR': {'POLICY': 'STOP'} }, @@ -240,8 +269,9 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f # Assert assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called_once_with( + gates_activity.debug.assert_called() + gates_activity.send_notification.assert_called_once_with( + metadata=metadata['metadata'], notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR", message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error", block="mlflow_gate", @@ -254,6 +284,7 @@ async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_f async def test_mlflow_content_gate_no_filters(gates_activity): # Arrange input_data = { + **metadata, 'filters': {}, 'data': {'value': [1, 2, 3]}, 'type': 'test', @@ -265,13 +296,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity): # Assert assert result == (None, 0, "") - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_mlflow_content_gate_with_filter(gates_activity): # Arrange input_data = { + **metadata, 'filters': { 'NAN_VALUES': {'policy': 'STOP', 'config': {}} }, @@ -286,14 +318,15 @@ async def test_mlflow_content_gate_with_filter(gates_activity): # Assert assert result == ( 'STOP', -1, "Transformed data not passed the content filter") - gates_activity.logger.debug.assert_called() - gates_activity.notification_handler.build_and_send_notification.assert_called() + gates_activity.debug.assert_called() + gates_activity.send_notification.assert_called() @mark.asyncio async def test_format_prediction(gates_activity): # Arrange input_data = { + **metadata, 'data': {'prediction': [1], 'response_time': [0.1]}, 'timestamp': '2023-05-26 11:12:27', 'model_id': 'test_model', @@ -311,13 +344,14 @@ async def test_format_prediction(gates_activity): assert result['prediction_confidence'] == {0: 0.9} assert result['prediction_status'] == {0: 'Good'} assert result['comments'] == {0: ""} - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_format_default_prediction(gates_activity): # Arrange input_data = { + **metadata, 'timestamp': '2023-05-26 11:12:27', 'model_id': 'test_model', 'prediction_confidence': 0.1, @@ -335,13 +369,14 @@ async def test_format_default_prediction(gates_activity): assert result['prediction_confidence'] == {0: 0.1} assert result['prediction_status'] == {0: 'Bad'} assert result['comments'] == {0: 'Test comment'} - gates_activity.logger.debug.assert_called() + gates_activity.debug.assert_called() @mark.asyncio async def test_get_last_timestamp_with_data(gates_activity): # Arrange input_data = { + **metadata, 'data': { 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'] } diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 53dd018..a92f3f3 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -39,6 +39,16 @@ def mlflow(mock_mlflow_repository): ) +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + @mark.asyncio @patch("laborious.activities.mlflow.DataFrame") @patch("laborious.activities.mlflow.max") @@ -46,6 +56,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { + **metadata, 'data': [ {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0, 'created_at': '2024-01-01 12:00:00'}, @@ -100,6 +111,7 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow): mock_max.return_value = '2024-01-02' # Mock input data input_data = { + **metadata, 'data': [ {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0}, {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0}, diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index bbdc626..ebce052 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -5,6 +5,15 @@ from laborious.activities.opc import NotificationLevel from laborious.activities.opc import OPC +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + @patch("laborious.activities.opc.OpcRepository") def test___init__(mock_opc_repository): @@ -15,6 +24,7 @@ def test___init__(mock_opc_repository): mock_notification_handler = MagicMock() servers = { 'server1': { + 'id': 'server1', 'url': 'http://localhost:8080', 'server_uri': 'opc.tcp://localhost:4840', 'cert_path': '', @@ -23,6 +33,7 @@ def test___init__(mock_opc_repository): 'reconnection_interval': 60, }, 'server2': { + 'id': 'server2', 'url': 'http://localhost:8080', 'server_uri': 'opc.tcp://localhost:4840', 'cert_path': '', @@ -45,7 +56,7 @@ def test___init__(mock_opc_repository): mock_opc_repository.assert_has_calls([ call( - name="server1", + id="server1", url="http://localhost:8080", logger=mock_logger, server_uri="opc.tcp://localhost:4840", @@ -58,7 +69,7 @@ def test___init__(mock_opc_repository): ]) mock_opc_repository.assert_has_calls([ call( - name="server2", + id="server2", url="http://localhost:8080", logger=mock_logger, server_uri="opc.tcp://localhost:4840", @@ -79,6 +90,7 @@ def test___init__(mock_opc_repository): def opc(mock_opc_repository): servers = { 'server1': { + 'id': 'server1', 'url': 'http://localhost:8080', 'server_uri': 'opc.tcp://localhost:4840', 'cert_path': '', @@ -90,12 +102,15 @@ def opc(mock_opc_repository): mock_opc_repository.write_data = MagicMock( return_value=True ) - return OPC( + opc = OPC( opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock() ) + opc.send_notification = MagicMock() + return opc + WRITE_DATA_CASES = [ ('tag1', 'int', 50), @@ -107,10 +122,10 @@ WRITE_DATA_CASES = [ @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES) def test_write_data_success(opc, tag, data_type, data): - assert opc.write_data(server='server1', tag=tag, data=data, - data_type=data_type, tag_type='prediction') + assert opc.write_data(server_id='server1', tag=tag, data=data, + data_type=data_type, tag_type='prediction', metadata=metadata) opc.opc_repository['server1'].write_data.assert_called_once_with( - tag, data, data_type) + tag, data, data_type, opc.logger, metadata) def test_write_data_exception(opc): @@ -118,11 +133,12 @@ def test_write_data_exception(opc): "Test error") try: - opc.write_data(server='server1', tag='tag1', data=50, - data_type='int', tag_type='prediction') + opc.write_data(server_id='server1', tag='tag1', data=50, + data_type='int', tag_type='prediction', metadata=metadata) except Exception: - opc.notification_handler.build_and_send_notification.assert_called_once_with( + opc.send_notification.assert_called_once_with( + metadata=metadata, notification_id="WRITE_OPC_PREDICTION_ERROR", message="Error writing data to OPC server: Test error", block="write_opc_data", @@ -138,6 +154,7 @@ def test_write_data_exception(opc): async def test_write_opc_data_success(opc): # Arrange input_data = { + **metadata, 'data': { 'prediction': [0.75], 'prediction_confidence': [0.95] @@ -163,19 +180,21 @@ async def test_write_opc_data_success(opc): assert output == {'data': 'data'} opc.write_data.assert_has_calls([ call( - server='server1', + server_id='server1', tag='tag1', data=0.75, data_type='float', - tag_type='prediction' + tag_type='prediction', + metadata=metadata['metadata'] )]) opc.write_data.assert_has_calls([ call( - server='server1', + server_id='server1', tag='tag2', data=0.95, data_type='float', - tag_type='confidence' + tag_type='confidence', + metadata=metadata['metadata'] ) ]) assert opc.write_data.call_count == 2 @@ -185,6 +204,7 @@ async def test_write_opc_data_success(opc): async def test_write_opc_data_empty_config(opc): # Arrange input_data = { + **metadata, 'data': { 'prediction': [0.75], 'prediction_confidence': [0.95] @@ -209,7 +229,7 @@ async def test_write_opc_data_empty_config(opc): ]) def test_process_confidence(opc, data, success, expected): # Act - result = opc.process_confidence(data, success) + result = opc.process_confidence(data, success, metadata) # Assert assert result['prediction_confidence'][0] == expected diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index 7369496..1b661d2 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -14,7 +14,7 @@ def mock_logger(): @fixture def opc_repository(mock_logger): return OpcRepository( - name="test_repo", + id="test_repo", url="opc.tcp://localhost:4840", logger=mock_logger, notification_handler=Mock(), @@ -34,8 +34,18 @@ def mock_client(): yield client_instance +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + def test_init(opc_repository): - assert opc_repository.name == "test_repo" + assert opc_repository.id == "test_repo" assert opc_repository.url == "opc.tcp://localhost:4840" assert opc_repository.server_uri == "urn:test:server" assert opc_repository.cert_path == "/path/to/cert.pem" @@ -110,7 +120,7 @@ def test_try_connect_fail(opc_repository): opc_repository.client.connect.assert_called_once() assert opc_repository.last_reconnection_time is not None opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.name}", + notification_id=f"OPC_CONNECTION_ERROR_{opc_repository.id}", message="Failed to connect to OPC server: Test error", block="opc_repository", level=NotificationLevel.ERROR, @@ -126,6 +136,17 @@ def test_disconnect(opc_repository, mock_client): assert opc_repository.client is None +def test_disconnect_error(opc_repository, mock_client): + opc_repository.client = mock_client + mock_client.disconnect.side_effect = Exception("Test error") + opc_repository.disconnect() + + opc_repository.logger.error.assert_called_once_with( + "Failed to disconnect from OPC server: Test error" + ) + assert opc_repository.client is None + + def test_validate_connection_none_client(opc_repository): opc_repository.client = None opc_repository.connect = MagicMock() @@ -174,11 +195,11 @@ def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repositor opc_repository.client = MagicMock() opc_repository.client.aio_obj.uaclient.protocol = None opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0) - opc_repository.try_connect = MagicMock() + opc_repository.connect = MagicMock() response = opc_repository.validate_connection() - opc_repository.try_connect.assert_called_once() - assert response == opc_repository.try_connect.return_value + opc_repository.connect.assert_called_once() + assert response == opc_repository.connect.return_value def test_validate_connection_failed(opc_repository): @@ -192,7 +213,8 @@ def test_validate_connection_failed(opc_repository): def test_write_data_validate_connection_do_nothing(opc_repository): opc_repository.validate_connection = MagicMock(return_value=True) opc_repository.client = MagicMock() - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") @@ -201,7 +223,8 @@ def test_write_data_validate_connection_failed(opc_repository): opc_repository.validate_connection = MagicMock(return_value=False) opc_repository.client = MagicMock() opc_repository.error_count = 0 - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_not_called() @@ -211,12 +234,13 @@ def test_write_data_get_node_failed(opc_repository): opc_repository.client = MagicMock() opc_repository.error_count = 0 opc_repository.client.get_node.side_effect = Exception("Test error") - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode") opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.name}", - message="Failed to get node from OPC server: Test error", + notification_id=f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}", + message="Failed to get node from OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}", block="opc_repository", level=NotificationLevel.ERROR, attachment_content=ANY @@ -230,13 +254,14 @@ def test_write_data_invalid_data_type(opc_repository, mock_client): mock_node = MagicMock() mock_client.get_node.return_value = mock_node - opc_repository.write_data("ns=2;s=TestNode", 42.0, "invalid_type") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "invalid_type", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.name}", - message="Unsupported data type: invalid_type", + notification_id=f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}", + message="Unsupported data type: invalid_type | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}", block="opc_repository", level=NotificationLevel.ERROR ) @@ -248,12 +273,11 @@ def test_write_data(opc_repository, mock_client): mock_node = MagicMock() mock_client.get_node.return_value = mock_node - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata) mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") mock_node.write_value.assert_called_once() - opc_repository.logger.info.assert_called_once_with( - "Writing 42.0 - to " + str(mock_node)) def test_write_data_write_value_failed(opc_repository, mock_client): @@ -263,13 +287,14 @@ def test_write_data_write_value_failed(opc_repository, mock_client): opc_repository.error_count = 0 mock_client.get_node.return_value = mock_node mock_node.write_value.side_effect = Exception("Test error") - opc_repository.write_data("ns=2;s=TestNode", 42.0, "float") + opc_repository.write_data("ns=2;s=TestNode", 42.0, + "float", opc_repository.logger, metadata) opc_repository.validate_connection.assert_called_once() mock_client.get_node.assert_called_once_with("ns=2;s=TestNode") mock_node.write_value.assert_called_once() opc_repository.notification_handler.build_and_send_notification.assert_called_once_with( - notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.name}", - message="Failed to write data to OPC server: Test error", + notification_id=f"OPC_WRITE_DATA_ERROR_{opc_repository.id}", + message="Failed to write data to OPC server: Test error | metadata: {'metadata': {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}}", block="opc_repository", level=NotificationLevel.ERROR, attachment_content=ANY diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index b137bc2..a7b36aa 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -54,7 +54,7 @@ def test_build_opc_config_with_env_vars(): def test_build_opc_config_with_individual_env_vars(): # Arrange environ.pop('OPC_CONFIG', None) - environ['OPC_NAME'] = 'test-name' + environ['OPC_ID'] = '1' environ['OPC_URL'] = 'opc.tcp://test:4840' environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840' environ['OPC_RECONNECTION_INTERVAL'] = '300' @@ -63,16 +63,16 @@ def test_build_opc_config_with_individual_env_vars(): config = build_opc_config() # Assert - assert config['opc']['name'] == 'test-name' - assert config['opc']['url'] == 'opc.tcp://test:4840' - assert config['opc']['server_uri'] == 'opc.tcp://test:4840' - assert config['opc']['reconnection_interval'] == 300 + assert config['1']['id'] == '1' + assert config['1']['url'] == 'opc.tcp://test:4840' + assert config['1']['server_uri'] == 'opc.tcp://test:4840' + assert config['1']['reconnection_interval'] == 300 def test_build_opc_config_with_defaults(): # Arrange environ.pop('OPC_CONFIG', None) - environ.pop('OPC_NAME', None) + environ.pop('OPC_ID', None) environ.pop('OPC_URL', None) environ.pop('OPC_SERVER_URI', None) environ.pop('OPC_RECONNECTION_INTERVAL', None) @@ -81,10 +81,10 @@ def test_build_opc_config_with_defaults(): config = build_opc_config() # Assert - assert config['opc']['name'] == 'opc' - assert config['opc']['url'] == 'opc.tcp://localhost:4840' - assert config['opc']['server_uri'] == 'opc.tcp://localhost:4840' - assert config['opc']['reconnection_interval'] == 120 + assert config['1']['id'] == '1' + assert config['1']['url'] == 'opc.tcp://localhost:4840' + assert config['1']['server_uri'] == 'opc.tcp://localhost:4840' + assert config['1']['reconnection_interval'] == 120 def test_build_postgres_config_with_env_vars(): diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index 85cbd34..7dbf39a 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -10,11 +10,22 @@ def format_and_export_prediction(): return FormatAndExportPrediction() +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + @mark.asyncio @patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock) async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): input_data = { + 'metadata': metadata, "path_flag": None, "data": {"test": "data"}, "timestamp": "2021-01-01", @@ -35,7 +46,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'data': input_data['data'], 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], - 'prediction_confidence': input_data['prediction_confidence'] + 'prediction_confidence': input_data['prediction_confidence'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY @@ -46,7 +58,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): Activities.write_opc_data, { 'opc_output_config': input_data['opc_output_config'], - 'data': workflow_mock.execute_local_activity_method.return_value + 'data': workflow_mock.execute_local_activity_method.return_value, + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY @@ -59,7 +72,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): { 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value + 'data': workflow_mock.execute_activity_method.return_value, + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY @@ -74,6 +88,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): async def test_run_default_path_flag(workflow_mock, format_and_export_prediction): input_data = { + 'metadata': metadata, "path_flag": "default", "data": {"test": "data"}, "timestamp": "2021-01-01", @@ -95,7 +110,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], - 'comment': input_data['comment'] + 'comment': input_data['comment'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY @@ -107,7 +123,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction Activities.write_opc_data, { 'opc_output_config': input_data['opc_output_config'], - 'data': workflow_mock.execute_local_activity_method.return_value + 'data': workflow_mock.execute_local_activity_method.return_value, + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY @@ -120,7 +137,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction { 'schema': input_data['schema'], 'table_name': input_data['table_name'], - 'data': workflow_mock.execute_activity_method.return_value + 'data': workflow_mock.execute_activity_method.return_value, + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 560565e..e76b586 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -9,12 +9,23 @@ def prediction_process(): return PredictionProcess() +metadata = { + "metadata": { + "model_id": "test_model", + "model_name": "test_model", + "workflow_name": "test_workflow", + "schema_name": "test_schedule", + }, +} + + @mark.asyncio @patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock) async def test_run(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=False) # Arrange input_data = { + 'metadata': metadata, 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', @@ -49,51 +60,61 @@ async def test_run(workflow_mock, prediction_process): assert workflow_mock.execute_local_activity_method.call_count == 7 workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, - retry_policy=ANY, start_to_close_timeout=ANY)]) + call(Activities.get_last_timestamp, { + 'data': input_data['data'], + **metadata + }, + retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.input_gate, { 'filters': input_data['input_filters'], 'data': input_data['data'], - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_transform, { 'data': input_data['data'], 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_retention': input_data['model_retention'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_response_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': 'transformed_data', 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_predict, { 'data': 'transformed_data', 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_retention': input_data['model_retention'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_response_gate, { 'filters': input_data['mlflow_predict_filters'], 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'type': 'predict', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_child_workflow.assert_called_once_with( 'format_and_export_prediction', { + 'metadata': metadata, 'path_flag': 'continue', 'data': 'predicted_data', 'prediction_confidence': 0.95, @@ -115,6 +136,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=True) # Arrange input_data = { + 'metadata': metadata, 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', @@ -141,11 +163,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): assert workflow_mock.execute_local_activity_method.call_count == 2 workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.get_last_timestamp, { - 'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY), + 'data': input_data['data'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY), call(Activities.input_gate, { 'filters': input_data['input_filters'], 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY) + 'path_priority': input_data['path_priority'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY) ]) workflow_mock.execute_child_workflow.assert_not_called() @@ -156,6 +182,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True]) # Arrange input_data = { + 'metadata': metadata, 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', @@ -183,19 +210,26 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ # Assert assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, + call(Activities.get_last_timestamp, { + 'data': input_data['data'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.input_gate, { 'filters': input_data['input_filters'], 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, + 'path_priority': input_data['path_priority'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_transform, { 'data': input_data['data'], 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention']}, + 'model_retention': input_data['model_retention'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY) ]) workflow_mock.execute_local_activity_method.assert_has_calls([ @@ -203,7 +237,8 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'filters': input_data['mlflow_transform_filters'], 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY) ]) workflow_mock.execute_child_workflow.assert_not_called() @@ -216,6 +251,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process side_effect=[False, False, True]) # Arrange input_data = { + 'metadata': metadata, 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', @@ -247,33 +283,41 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process assert workflow_mock.execute_local_activity_method.call_count == 5 workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, + call(Activities.get_last_timestamp, { + 'data': input_data['data'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.input_gate, { 'filters': input_data['input_filters'], 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, + 'path_priority': input_data['path_priority'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_transform, { 'data': input_data['data'], 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_retention': input_data['model_retention'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_response_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': 'transformed_data', 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_child_workflow.assert_not_called() @@ -285,6 +329,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p side_effect=[False, False, False, True]) # Arrange input_data = { + 'metadata': metadata, 'data': {'test': 'data'}, 'schema': 'test_schema', 'table_name': 'test_table', @@ -317,46 +362,56 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p # Assert assert workflow_mock.execute_local_activity_method.call_count == 7 workflow_mock.execute_local_activity_method.assert_has_calls([ - call(Activities.get_last_timestamp, {'data': input_data['data']}, + call(Activities.get_last_timestamp, { + 'data': input_data['data'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.input_gate, { 'filters': input_data['input_filters'], 'data': input_data['data'], - 'path_priority': input_data['path_priority']}, + 'path_priority': input_data['path_priority'], + **metadata + }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_transform, { 'data': input_data['data'], 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_retention': input_data['model_retention'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_response_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], 'data': 'transformed_data', 'type': 'transform', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_predict, { 'data': 'transformed_data', 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] + 'model_retention': input_data['model_retention'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_response_gate, { 'filters': input_data['mlflow_predict_filters'], 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'type': 'predict', - 'path_priority': input_data['path_priority'] + 'path_priority': input_data['path_priority'], + **metadata }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_child_workflow.assert_not_called() @@ -378,6 +433,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process): # Act result = await prediction_process.path_flag_handler( data, path_flag, { + 'metadata': metadata, 'schema': schema, 'table_name': table_name, 'model_id': model, @@ -410,6 +466,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): # Act result = await prediction_process.path_flag_handler( data, path_flag, { + 'metadata': metadata, 'schema': schema, 'table_name': table_name, 'model_id': model, @@ -424,6 +481,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process): workflow_mock.execute_activity_method.assert_called_once_with( Activities.repeat_last_prediction, { + **metadata, 'schema': schema, 'table_name': table_name, 'model': model, @@ -452,6 +510,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): # Act result = await prediction_process.path_flag_handler( data, path_flag, { + 'metadata': metadata, 'schema': schema, 'table_name': table_name, 'model_id': model, @@ -468,6 +527,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): workflow_mock.execute_child_workflow.assert_called_once_with( 'format_and_export_prediction', { + 'metadata': metadata, 'path_flag': path_flag, 'data': data, 'prediction_confidence': confidence, @@ -500,6 +560,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): # Act result = await prediction_process.path_flag_handler( data, path_flag, { + **metadata, 'schema': schema, 'table_name': table_name, 'model_id': model, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 0ca45e1..3c8fc82 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -9,6 +9,16 @@ def predictions_batch() -> PredictionsBatch: return PredictionsBatch() +metadata = { + "metadata": { + "model_id": "test_model_id", + "model_name": "test_model", + "workflow_name": "predictions_batch", + "schedule_name": "test_schedule", + }, +} + + @mark.asyncio @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): @@ -29,27 +39,17 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch workflow_mock.execute_local_activity_method.assert_has_calls([ call( - Activities.prepare_activity, + Activities.load_custom_query, { - 'schedule_name': input_data['schedule_name'], - 'model_name': input_data['model_name'], - 'model_id': input_data['model_id'], - 'workflow_name': 'predictions_batch' + **metadata, + 'query': input_data['query'], }, retry_policy=ANY, start_to_close_timeout=ANY ) ]) - - workflow_mock.execute_local_activity_method.assert_has_calls([ - call( - Activities.load_custom_query, - input_data['query'], - retry_policy=ANY, - start_to_close_timeout=ANY - ) - ]) prediction_input = { + 'metadata': metadata, 'data': {'data': 'test_data'}, 'schema': input_data['schema'], 'table_name': input_data['table_name'],