From cfb217963b5a7d40327028f7cb91c5317254f2b3 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 28 May 2025 09:24:16 -0300 Subject: [PATCH] SIENTIAPDE-994 Add shutdown tests for Activities and OPC classes; improve test assertions --- tests/laborious/activities/test_activities.py | 44 +++++++++++++++++++ tests/laborious/activities/test_opc.py | 5 +++ .../utils/repository/test_model_repository.py | 4 +- .../subworkflows/test_prediction_process.py | 15 ++++--- 4 files changed, 60 insertions(+), 8 deletions(-) diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 98be011..3b2ef49 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -147,3 +147,47 @@ async def test_prepare_activity(_mock_opc_init, '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()) +def test_shutdown(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 + ) + + activities.shutdown() + mock_opc_init.shutdown.assert_called_once() + mock_postgres_init.close.assert_called_once() diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index c4e26ee..d012778 100644 --- a/tests/laborious/activities/test_opc.py +++ b/tests/laborious/activities/test_opc.py @@ -189,3 +189,8 @@ async def test_write_opc_data_empty_config(opc): # Assert opc.opc_repository['server1'].write_data.assert_not_called() + + +def test_shutdown(opc): + opc.shutdown() + opc.opc_repository['server1'].disconnect.assert_called_once() diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index abf4ab8..a675edb 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -251,9 +251,9 @@ def test_predict_success(mlflow_repository): mlflow_repository.model_serving.get_cached_predict.assert_called_once_with( model_name, data, 1) - assert output['success'] == True + assert output['success'] is True assert output['content'] == {'prediction': { - 0: 2, 1: 3}, 'response_time': ANY} + 0: 3}, 'response_time': ANY} def test_predict_error(mlflow_repository): diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 4aebc6f..4318379 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -73,13 +73,13 @@ async def test_run(workflow_mock, prediction_process): workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'data': 'transformed_data', 'type': 'transform', 'path_priority': input_data['path_priority'] }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_predict, { - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'data': 'transformed_data', 'model_name': input_data['model_name'], 'model_retention': input_data['model_retention'] }, retry_policy=ANY, start_to_close_timeout=ANY)]) @@ -101,7 +101,10 @@ async def test_run(workflow_mock, prediction_process): 'model_id': 1, 'model_name': 'test_model_name', 'model_retention': '30', - 'opc_output_config': input_data['opc_output_config'] + 'opc_output_config': input_data['opc_output_config'], + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'comment': 'Error' } ) @@ -268,7 +271,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'data': 'transformed_data', 'type': 'transform', 'path_priority': input_data['path_priority'] }, retry_policy=ANY, start_to_close_timeout=ANY)]) @@ -338,13 +341,13 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'data': 'transformed_data', 'type': 'transform', 'path_priority': input_data['path_priority'] }, retry_policy=ANY, start_to_close_timeout=ANY)]) workflow_mock.execute_local_activity_method.assert_has_calls([ call(Activities.request_predict, { - 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + 'data': 'transformed_data', 'model_name': input_data['model_name'], 'model_retention': input_data['model_retention'] }, retry_policy=ANY, start_to_close_timeout=ANY)])