SIENTIAPDE-994

Add shutdown tests for Activities and OPC classes; improve test assertions

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

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

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

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