SIENTIAPDE-994

Refactor tests for Postgres activities and improve error handling

- Updated test_postgres.py to enhance the testing of load_custom_query method, including cases for None data and date conversion.
- Refactored repeat_last_prediction tests to use mocks for SQLAlchemy session execution.
- Added tests for export_data_to_postgres method, covering both success and error scenarios.
- Improved the initialization tests for Activities class to ensure proper instantiation of dependencies.
- Enhanced test coverage for OPC repository connection validation.
- Updated tests for prediction workflows to streamline input handling and improve clarity.
- Introduced tests for connectors configuration to validate environment variable handling for MLFlow, OPC, and Postgres.
- Added tests for logger utility to ensure default settings are correctly applied.

.gitignore

@@ -32,4 +32,8 @@ __pycache__/
 *.tmp
 *.bak
 *.old
-.secret
+.secret
+
+# Ignorar coverage
+htmlcov/
+.coverage

laborious/activities/gates.py

@@ -137,15 +137,25 @@ class Gates(BaseActivity):
         for fil, config in filters.items():
             if fil not in mlflow_response_filter_functions:
                 continue
-            if mlflow_response_filter_functions[fil](data, config):
+            try:
-                filter_output.append(config['POLICY'])
+                if mlflow_response_filter_functions[fil](data, config):
-                comments.append(data['content']['message'])
+                    filter_output.append(config['POLICY'])
+                    comments.append(data['content']['message'])
+                    self.notification_handler.build_and_send_notification(
+                        notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
+                        message=data['content']['message'],
+                        block="mlflow_gate",
+                        level=NotificationLevel.WARNING,
+                        attachment_content=data['content']['traceback']
+                    )
+            except Exception as e:
+                trace = traceback.format_exc()
                 self.notification_handler.build_and_send_notification(
-                    notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
+                    notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
-                    message=data['content']['message'],
+                    message=f"Error in filter {fil}:{config}: \n {e}",
                     block="mlflow_gate",
-                    level=NotificationLevel.WARNING,
+                    level=NotificationLevel.ERROR,
-                    attachment_content=data['content']['traceback']
+                    attachment_content=trace
                 )
 
         for path_flag in path_priority:
@@ -189,14 +199,24 @@ class Gates(BaseActivity):
         for fil, config in filters.items():
             if fil not in mlflow_content_filter_functions:
                 continue
-            if mlflow_content_filter_functions[fil](data, config):
+            try:
-                filter_output.append(config['POLICY'])
+                if mlflow_content_filter_functions[fil](data, config):
+                    filter_output.append(config['POLICY'])
+                    self.notification_handler.build_and_send_notification(
+                        notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
+                        message=f"Data not passed the content filter {fil}:{config}",
+                        block="mlflow_gate",
+                        level=NotificationLevel.WARNING,
+                        attachment_content=data.to_string()
+                    )
+            except Exception as e:
+                trace = traceback.format_exc()
                 self.notification_handler.build_and_send_notification(
-                    notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
+                    notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
-                    message=f"Data not passed the content filter {fil}:{config}",
+                    message=f"Error in filter {fil}:{config}: \n {e}",
                     block="mlflow_gate",
-                    level=NotificationLevel.WARNING,
+                    level=NotificationLevel.ERROR,
-                    attachment_content=data.to_string()
+                    attachment_content=trace
                 )
 
         for path_flag in path_priority:

laborious/utils/filters/mlflow_filters.py

@@ -14,7 +14,7 @@ def api_error_filter(response: dict, _config: dict):
 
 def nan_values_filter(predictions: DataFrame, _config: dict):
     data = predictions.replace({None: np.nan}).drop(
-        columns=['timestamp'], errors='ignore')
+        columns=['timestamp'], errors='ignore').infer_objects(copy=False)
 
     if data.isna().all().all():
         return True

laborious/utils/repository/opc_repository.py

@@ -122,12 +122,17 @@ class OpcRepository():
             return False
 
     def disconnect(self):
+        if self.client is None:
+            return
         self.client.disconnect()
         self.client = None
         self.logger.info('Disconnected from OPC server')
 
     def __del__(self):
-        self.disconnect()
+        try:
+            self.disconnect()
+        except Exception as e:
+            self.logger.error(f"Error in destructor: {e}")
 
     def validate_connection(self):
         if self.client is None:

tests/laborious/activities/test_activities.py

@@ -0,0 +1,149 @@
+from pytest import mark
+from unittest.mock import patch, MagicMock, ANY
+from laborious.activities.activities import Activities
+from laborious.activities.postgres import Postgres
+from laborious.activities.mlflow import MLFlow
+from laborious.activities.gates import Gates
+from laborious.activities.opc import OPC
+
+
+@patch('laborious.activities.activities.Postgres.__init__')
+@patch('laborious.activities.activities.MLFlow.__init__')
+@patch('laborious.activities.activities.OPC.__init__')
+@patch('laborious.activities.activities.Gates.__init__')
+def test___init__(mock_gates_init, 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
+    )
+
+    assert isinstance(activities, Activities)
+    assert isinstance(activities, Postgres)
+    assert isinstance(activities, MLFlow)
+    assert isinstance(activities, OPC)
+    assert isinstance(activities, Gates)
+
+    mock_postgres_init.assert_called_once_with(
+        ANY,
+        host=postgres_config['host'],
+        port=postgres_config['port'],
+        user=postgres_config['user'],
+        password=postgres_config['password'],
+        dbname=postgres_config['dbname'],
+        min_connections=postgres_config['min_connections'],
+        max_connections=postgres_config['max_connections'],
+        logger=logger,
+        notification_handler=notification_handler
+    )
+
+    mock_mlflow_init.assert_called_once_with(
+        ANY,
+        mlflow_host=mlflow_config['host'],
+        mlflow_port=mlflow_config['port'],
+        mlflow_username=mlflow_config['username'],
+        mlflow_password=mlflow_config['password'],
+        logger=logger,
+        notification_handler=notification_handler
+    )
+
+    mock_opc_init.assert_called_once_with(
+        ANY,
+        opc_servers=opc_config,
+        logger=logger,
+        notification_handler=notification_handler
+    )
+
+    mock_gates_init.assert_called_once_with(
+        ANY,
+        logger=logger,
+        notification_handler=notification_handler
+    )
+
+
+@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_name == input_data[
+        'workflow_name']
+    assert activities.notification_handler.base_notification.schedule_name == 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']

tests/laborious/activities/test_gates.py

@@ -1,605 +1,369 @@
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import MagicMock, ANY, patch
-from pandas import DataFrame
 from pytest import fixture, mark
-
-from laborious.activities.gates import Gates
 from sientia_do.notifications.models import NotificationLevel
+from laborious.activities.gates import Gates
 
 
 @fixture
-def gates():
+def gates_activity():
     return Gates(
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
+    )
+
+
+@mark.asyncio
+async def test_input_gate_invalid_filter(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {
+            'INVALID_FILTER': {'POLICY': 'STOP'}
+        },
+        'data': {'value': [1, 2, 3]},
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.input_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
+    gates_activity.logger.error.assert_called_once_with(
+        "Filter INVALID_FILTER not found"
     )
 
 
 @mark.asyncio
 @patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_specific_variables_null_values_with_stop_policy_only(
+async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
-    input_filter_functions_mock,
+    # Arrange
-    gates
+    mock_input_filter_functions.__contains__.return_value = True
-):
+    mock_input_filter_functions.__getitem__.return_value = MagicMock(
-    specific_variables_null_values_mock = MagicMock(return_value=True)
+        side_effect=Exception("Test error"))
-    empty_data_mock = MagicMock(return_value=False)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
     input_data = {
         'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
+            'EMPTY_DATA': {'POLICY': 'STOP'}
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            }
         },
-        'data': {
+        'data': {'value': []},
-            'variable': ['variable1', 'variable2'],
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
     }
 
-    result = await gates.input_gate(input_data)
+    # Act
-    assert result == ('stop', -1, 'Input data with bad quality')
+    result = await gates_activity.input_gate(input_data)
 
-    input_args = specific_variables_null_values_mock.call_args
+    # Assert
-    assert input_args[0][0].equals(DataFrame(
+    assert result == (None, 0, "")
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
+    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
-    assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
+        notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
-
+        message="Error in filter EMPTY_DATA:{'POLICY': 'STOP'}: \n Test error",
-    empty_data_mock.assert_not_called()
+        block="input_gate",
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_specific_variables_null_values_with_continue_policy_only(
-    input_filter_functions_mock,
-    gates
-):
-    specific_variables_null_values_mock = MagicMock(return_value=True)
-    empty_data_mock = MagicMock(return_value=False)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'continue',
-                'VARIABLES': ['variable2']
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == ('continue', 2, 'Input data with bad quality')
-
-    input_args = specific_variables_null_values_mock.call_args
-    assert input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert input_args[0][1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
-
-    empty_data_mock.assert_not_called()
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_specific_variables_null_values_no_filtered(
-    input_filter_functions_mock,
-    gates
-):
-    specific_variables_null_values_mock = MagicMock(return_value=False)
-    empty_data_mock = MagicMock(return_value=False)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == (None, 0, '')
-
-    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
-
-    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert specific_variables_null_values_input_args[0][
-        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
-
-    empty_data_mock.assert_not_called()
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_one_stop_policy(
-    input_filter_functions_mock,
-    gates
-):
-    specific_variables_null_values_mock = MagicMock(return_value=True)
-    empty_data_mock = MagicMock(return_value=True)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            },
-            'EMPTY_DATA': {
-                'POLICY': 'continue',
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == ('stop', -1, 'Input data with bad quality')
-
-    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
-    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert specific_variables_null_values_input_args[0][
-        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
-
-    empty_data_input_args = empty_data_mock.call_args
-    assert empty_data_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_one_continue_policy(
-    input_filter_functions_mock,
-    gates
-):
-    specific_variables_null_values_mock = MagicMock(return_value=False)
-    empty_data_mock = MagicMock(return_value=True)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            },
-            'EMPTY_DATA': {
-                'POLICY': 'continue',
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == ('continue', 2, 'Input data with bad quality')
-
-    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
-    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert specific_variables_null_values_input_args[0][
-        1] == input_data['filters']['SPECIFIC_VARIABLES_NULL_VALUES']
-
-    empty_data_input_args = empty_data_mock.call_args
-    assert empty_data_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_no_filtered(
-    input_filter_functions_mock,
-    gates
-):
-    specific_variables_null_values_mock = MagicMock(return_value=False)
-    empty_data_mock = MagicMock(return_value=False)
-
-    def functions_side_effect(x):
-        if x == 'SPECIFIC_VARIABLES_NULL_VALUES':
-            return specific_variables_null_values_mock
-        if x == 'path_confidence':
-            return {
-                'stop': -1,
-                'continue': 2,
-                'repeat': -1
-            }
-        return empty_data_mock
-
-    input_filter_functions_mock.__getitem__.side_effect = functions_side_effect
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            },
-            'EMPTY_DATA': {
-                'POLICY': 'continue',
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat']
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == (None, 0, '')
-
-    specific_variables_null_values_input_args = specific_variables_null_values_mock.call_args
-    assert specific_variables_null_values_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-
-    empty_data_input_args = empty_data_mock.call_args
-    assert empty_data_input_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert empty_data_input_args[0][1] == input_data['filters']['EMPTY_DATA']
-
-
-@mark.asyncio
-@patch('laborious.activities.gates.input_filter_functions')
-async def test_input_gate_error(
-    input_filter_functions_mock,
-    gates
-):
-    input_filter_functions_mock.__getitem__.side_effect = KeyError('test')
-
-    input_data = {
-        'filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'POLICY': 'stop',
-                'VARIABLES': ['variable2']
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat'],
-    }
-
-    result = await gates.input_gate(input_data)
-    assert result == (None, 0, '')
-
-    gates.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id='INTPUT_GATE_ERROR__SPECIFIC_VARIABLES_NULL_VALUES',
-        message="Error in filter SPECIFIC_VARIABLES_NULL_VALUES:{'POLICY': 'stop', 'VARIABLES': ['variable2']}: \n 'test'",
-        block='input_gate',
         level=NotificationLevel.ERROR,
         attachment_content=ANY
     )
 
 
-transform_filter_path_confidence = {
-    'stop': -1,
-    'continue': 255,
-    'repeat': -1
-}
-
-
 @mark.asyncio
-@patch('laborious.activities.gates.mlflow_response_filter_functions')
+async def test_input_gate_no_filters(gates_activity):
-async def test_mlflow_response_gate_no_filtered(
+    # Arrange
-    mlflow_response_filter_functions_mock,
-    gates
-):
-    api_error_filter_mock = MagicMock(return_value=False)
-
-    def transform_filter_functions_side_effect(x: str):
-        if x == 'API_ERROR':
-            return api_error_filter_mock
-        if x == 'path_confidence':
-            return transform_filter_path_confidence
-
-    mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
-
     input_data = {
-        'filters': {
+        'filters': {},
-            'API_ERROR': {
+        'data': {'value': [1, 2, 3]},
-                'POLICY': 'stop',
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
-            }
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat'],
-        'type': 'predict'
     }
 
-    result = await gates.mlflow_response_gate(input_data)
+    # Act
-    assert result == (None, 0, '')
+    result = await gates_activity.input_gate(input_data)
 
-    api_error_filter_mock.assert_called_once_with(
+    # Assert
-        input_data['data'],
+    assert result == (None, 0, "")
-        input_data['filters']['API_ERROR']
+    gates_activity.logger.debug.assert_called()
-    )
 
-    gates.notification_handler.build_and_send_notification.assert_not_called()
+
+@mark.asyncio
+async def test_input_gate_with_filter(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {
+            'EMPTY_DATA': {'POLICY': 'STOP'}
+        },
+        'data': {'value': []},
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.input_gate(input_data)
+
+    # Assert
+    assert result == ('STOP', -1, "Input data with bad quality")
+    gates_activity.logger.debug.assert_called()
+
+
+@mark.asyncio
+async def test_mlflow_response_gate_invalid_filter(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {
+            'INVALID_FILTER': {'POLICY': 'STOP'}
+        },
+        'data': {'content': {'message': 'success'}},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.mlflow_response_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
 
 
 @mark.asyncio
 @patch('laborious.activities.gates.mlflow_response_filter_functions')
-async def test_mlflow_response_gate_filtered(
+async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
-    mlflow_response_filter_functions_mock,
+                                                     gates_activity):
-    gates
+    # Arrange
-):
+    mock_mlflow_response_filter_functions.__contains__.return_value = True
-    api_error_filter_mock = MagicMock(return_value=True)
+    mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
-
+        side_effect=Exception("Test error"))
-    def transform_filter_functions_side_effect(x: str):
-        if x == 'API_ERROR':
-            return api_error_filter_mock
-        if x == 'path_confidence':
-            return transform_filter_path_confidence
-
-    mlflow_response_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
-
     input_data = {
         'filters': {
-            'API_ERROR': {
+            'INVALID_FILTER': {'POLICY': 'STOP'}
-                'POLICY': 'continue',
+        },
-            }
+        'data': {'content': {'message': 'success'}},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.mlflow_response_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
+    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
+        notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
+        message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
+        block="mlflow_gate",
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY
+    )
+
+
+@mark.asyncio
+async def test_mlflow_response_gate_no_filters(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {},
+        'data': {'content': {'message': 'success'}},
+        'type': 'test',
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.mlflow_response_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
+    gates_activity.logger.debug.assert_called()
+
+
+@mark.asyncio
+async def test_mlflow_response_gate_with_filter(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {
+            'API_ERROR': {'POLICY': 'STOP'}
         },
         'data': {
             'success': False,
             'content': {
-                'message': 'Error',
+                'message': 'API error occurred',
-                'traceback': 'Error'
+                'traceback': 'error trace'
             }
         },
-        'path_priority': ['stop', 'continue', 'repeat'],
+        'type': 'test',
-        'type': 'predict'
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
     }
 
-    result = await gates.mlflow_response_gate(input_data)
+    # Act
-    assert result == ('continue', 255, "Error")
+    result = await gates_activity.mlflow_response_gate(input_data)
 
-    api_error_filter_mock.assert_called_once_with(
+    # Assert
-        input_data['data'],
+    assert result == ('STOP', -1, "API error occurred")
-        input_data['filters']['API_ERROR']
+    gates_activity.logger.debug.assert_called()
-    )
+    gates_activity.notification_handler.build_and_send_notification.assert_called()
 
-    gates.notification_handler.build_and_send_notification.assert_called_once_with(
+
-        notification_id='PREDICT_GATE_RESPONSE_FILTER__API_ERROR',
+@mark.asyncio
-        message=input_data['data']['content']['message'],
+async def test_mlflow_content_gate_invalid_filter(gates_activity):
-        block='mlflow_gate',
+    # Arrange
-        level=NotificationLevel.WARNING,
+    input_data = {
-        attachment_content=input_data['data']['content']['traceback']
+        'filters': {
-    )
+            'INVALID_FILTER': {'POLICY': 'STOP'}
+        },
+        'data': {'value': [1, 2, 3]},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.mlflow_content_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
 
 
 @mark.asyncio
 @patch('laborious.activities.gates.mlflow_content_filter_functions')
-async def test_mlflow_content_gate_no_filtered(
+async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
-    mlflow_content_filter_functions_mock,
+                                                    gates_activity):
-    gates
+    # Arrange
-):
+    mock_mlflow_content_filter_functions.__contains__.return_value = True
-    nan_values_filter_mock = MagicMock(return_value=False)
+    mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
-
+        side_effect=Exception("Test error"))
-    def transform_filter_functions_side_effect(x: str):
-        if x == 'NAN_VALUES':
-            return nan_values_filter_mock
-        if x == 'path_confidence':
-            return transform_filter_path_confidence
-
-    mlflow_content_filter_functions_mock.__getitem__.side_effect = transform_filter_functions_side_effect
-
     input_data = {
         'filters': {
-            'NAN_VALUES': {
+            'API_ERROR': {'POLICY': 'STOP'}
-                'POLICY': 'repeat',
-            }
         },
         'data': {
-            'variable': ['variable1', 'variable2'],
+            'success': False,
-            'value': [1, 2]
+            'content': {
+                'message': 'API error occurred',
+                'traceback': 'error trace'
+            }
         },
-        'path_priority': ['stop', 'continue', 'repeat'],
+        'type': 'test',
-        'type': 'predict'
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
     }
 
-    result = await gates.mlflow_content_gate(input_data)
+    # Act
-    assert result == (None, 0, '')
+    result = await gates_activity.mlflow_content_gate(input_data)
 
-    nan_values_filter_mock_args = nan_values_filter_mock.call_args
+    # Assert
-    assert nan_values_filter_mock_args[0][0].equals(DataFrame(
+    assert result == (None, 0, "")
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
+    gates_activity.logger.debug.assert_called()
-    assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']
+    gates_activity.notification_handler.build_and_send_notification.assert_called_once_with(
-
+        notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
-    gates.notification_handler.build_and_send_notification.assert_not_called()
+        message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
+        block="mlflow_gate",
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY
+    )
 
 
 @mark.asyncio
-@patch('laborious.activities.gates.mlflow_content_filter_functions')
+async def test_mlflow_content_gate_no_filters(gates_activity):
-async def test_mlflow_content_gate_filtered(
+    # Arrange
-    mlflow_content_filter_functions_mock,
-    gates
-):
-    nan_values_filter_mock = MagicMock(return_value=True)
-
-    def transform_filter_functions_side_effect(x: str):
-        if x == 'NAN_VALUES':
-            return nan_values_filter_mock
-        if x == 'path_confidence':
-            return transform_filter_path_confidence
-
-    mlflow_content_filter_functions_mock.__getitem__.side_effect = \
-        transform_filter_functions_side_effect
-
     input_data = {
-        'filters': {
+        'filters': {},
-            'NAN_VALUES': {
+        'data': {'value': [1, 2, 3]},
-                'POLICY': 'repeat',
+        'type': 'test',
-            }
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
-        },
-        'data': {
-            'variable': ['variable1', 'variable2'],
-            'value': [1, 2]
-        },
-        'path_priority': ['stop', 'continue', 'repeat'],
-        'type': 'predict'
     }
 
-    result = await gates.mlflow_content_gate(input_data)
+    # Act
+    result = await gates_activity.mlflow_content_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, "")
+    gates_activity.logger.debug.assert_called()
+
+
+@mark.asyncio
+async def test_mlflow_content_gate_with_filter(gates_activity):
+    # Arrange
+    input_data = {
+        'filters': {
+            'NAN_VALUES': {'POLICY': 'STOP'}
+        },
+        'data': {'value': [None, None, None]},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+    }
+
+    # Act
+    result = await gates_activity.mlflow_content_gate(input_data)
+
+    # Assert
     assert result == (
-        'repeat', -1, "Transformed data not passed the content filter")
+        'STOP', -1, "Transformed data not passed the content filter")
-
+    gates_activity.logger.debug.assert_called()
-    nan_values_filter_mock_args = nan_values_filter_mock.call_args
+    gates_activity.notification_handler.build_and_send_notification.assert_called()
-    assert nan_values_filter_mock_args[0][0].equals(DataFrame(
-        {'variable': ['variable1', 'variable2'], 'value': [1, 2]}))
-    assert nan_values_filter_mock_args[0][1] == input_data['filters']['NAN_VALUES']
-
-    gates.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id='PREDICT_GATE_CONTENT_FILTER__NAN_VALUES',
-        message="Data not passed the content filter NAN_VALUES:{'POLICY': 'repeat'}",
-        block='mlflow_gate',
-        level=NotificationLevel.WARNING,
-        attachment_content=DataFrame(input_data['data']).to_string()
-    )
 
 
 @mark.asyncio
-async def test_format_prediction(
+async def test_format_prediction(gates_activity):
-    gates
+    # Arrange
-):
     input_data = {
-        'data': {
+        'data': {'prediction': [1], 'response_time': [0.1]},
-            'variable': ['variable1', 'variable2'],
+        'timestamp': '2023-05-26 11:12:27',
-            'value': [1, 2]
+        'model_id': 'test_model',
-        },
+        'prediction_confidence': 0.9
-        'timestamp': '2021-01-01',
-        'model_id': 'model_id',
-        'prediction_confidence': 0.95
     }
 
-    expected_output = DataFrame(input_data['data'])
+    # Act
-    expected_output['timestamp'] = input_data['timestamp']
+    result = await gates_activity.format_prediction(input_data)
-    expected_output['model_id'] = input_data['model_id']
-    expected_output['prediction_confidence'] = input_data['prediction_confidence']
-    expected_output['prediction_status'] = 'Good'
-    expected_output['comment'] = ''
 
-    result = await gates.format_prediction(input_data)
+    # Assert
-    assert result == expected_output.to_dict()
+    assert result['prediction'] == {0: 1}
+    assert result['response_time'] == {0: ANY}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['model_id'] == {0: 'test_model'}
+    assert result['prediction_confidence'] == {0: 0.9}
+    assert result['prediction_status'] == {0: 'Good'}
+    assert result['comments'] == {0: ""}
+    gates_activity.logger.debug.assert_called()
 
 
 @mark.asyncio
-async def test_format_default_prediction(
+async def test_format_default_prediction(gates_activity):
-    gates
+    # Arrange
-):
     input_data = {
-        'timestamp': '2021-01-01',
+        'timestamp': '2023-05-26 11:12:27',
-        'model_id': 'model_id',
+        'model_id': 'test_model',
-        'prediction_confidence': 0.95,
+        'prediction_confidence': 0.1,
-        'comment': 'Comment'
+        'comment': 'Test comment'
     }
 
-    expected_output = DataFrame({
+    # Act
-        'prediction': [0],
+    result = await gates_activity.format_default_prediction(input_data)
-        'response_time': [0],
-        'timestamp': [input_data['timestamp']],
-        'model_id': [input_data['model_id']],
-        'prediction_confidence': [input_data['prediction_confidence']],
-        'prediction_status': ['Bad'],
-        'comment': [input_data['comment']]
-    })
 
-    result = await gates.format_default_prediction(input_data)
+    # Assert
-    assert result == expected_output.to_dict()
+    assert result['prediction'] == {0: 0}
+    assert result['response_time'] == {0: 0}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['model_id'] == {0: 'test_model'}
+    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()
 
 
 @mark.asyncio
-async def test_get_last_timestamp(
+async def test_get_last_timestamp_with_data(gates_activity):
-    gates
+    # Arrange
-):
     input_data = {
         'data': {
-            'variable': ['variable1', 'variable2'],
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
-            'value': [1, 2],
-            'timestamp': ['2021-01-01', '2021-01-02']
         }
     }
 
-    result = await gates.get_last_timestamp(input_data)
+    # Act
-    assert result == '2021-01-02'
+    result = await gates_activity.get_last_timestamp(input_data)
+
+    # Assert
+    assert result == '2023-05-26 11:12:28'
+
+
+@mark.asyncio
+async def test_get_last_timestamp_no_data(gates_activity):
+    # Arrange
+    input_data = {
+        'data': {}
+    }
+
+    # Act
+    result = await gates_activity.get_last_timestamp(input_data)
+
+    # Assert
+    assert isinstance(result, str)  # Should be a timestamp string
+    assert len(result) > 0

tests/laborious/activities/test_postgres.py

@@ -1,132 +1,159 @@
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import MagicMock, patch
-from pandas import DataFrame
+from pytest import fixture, mark
-from pytest import fixture
+import pandas as pd
-from pytest import mark
-from sientia_do.notifications.models import NotificationLevel
-
 from laborious.activities.postgres import Postgres
 
 
 @fixture
-@patch("laborious.activities.postgres.ThreadedConnectionPool")
+@patch("laborious.activities.postgres.create_engine")
-def postgres_client(mock_pool):
+def postgres_activity(_mock_create_engine):
     return Postgres(
         host="localhost",
         port=5432,
-        user="postgres",
+        user="test_user",
-        password="postgres",
+        password="test_password",
-        dbname="postgres",
+        dbname="test_db",
         min_connections=1,
-        max_connections=10,
+        max_connections=5,
         logger=MagicMock(),
-        notification_handler=MagicMock(),
+        notification_handler=MagicMock()
     )
 
 
 @mark.asyncio
-@patch("laborious.activities.postgres.read_sql_query",
+@patch("laborious.activities.postgres.read_sql_query")
-       return_value=DataFrame([{"a": 1, "b": 2}]))
+async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
-async def test_load_custom_query_success(mock_read_sql_query, postgres_client):
+    query = "SELECT * FROM test_table LIMIT 1"
-    query = "SELECT * FROM test"
+    mock_read_sql_query.return_value = None
-    result = await postgres_client.load_custom_query(query)
+
-    assert result is not None
+    result = await postgres_activity.load_custom_query(query)
-    assert len(result) > 0
+
-    assert result == {'a': {0: 1}, 'b': {0: 2}}
+    assert isinstance(result, dict)
-    postgres_client.notification_handler.build_and_send_notification.assert_not_called()
+    assert len(result) == 0
 
 
 @mark.asyncio
-@patch("laborious.activities.postgres.read_sql_query",
+@patch("laborious.activities.postgres.read_sql_query")
-       side_effect=Exception("Error fetching data from query"))
+async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
-async def test_load_custom_query_error(mock_read_sql_query, postgres_client):
+    query = "SELECT * FROM test_table LIMIT 1"
-    query = "SELECT * FROM test"
+    mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
-    result = await postgres_client.load_custom_query(query)
+    mock_data['date'] = pd.to_datetime('2022-01-01')
-    assert result == {}
+
-    postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
+    mock_read_sql_query.return_value = mock_data
-        notification_id="ERROR_LOADING_CUSTOM_QUERY",
+
-        message="Error fetching data from query: Error fetching data from query",
+    result = await postgres_activity.load_custom_query(query)
-        block="load_custom_query",
+
-        level=NotificationLevel.ERROR,
+    assert isinstance(result, dict)
-        attachment_content=ANY
+    assert len(result) == 3
-    )
+    assert "column1" in result
+    assert "column2" in result
+    assert "date" in result
+    assert result['date'] == {0: '2022-01-01 00:00:00'}
 
 
 @mark.asyncio
-async def test_repeat_last_prediction_success(postgres_client):
+@patch("laborious.activities.postgres.read_sql_query")
-    query_items = {"schema": "test", "table_name": "test", "model": 1}
+async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
-    await postgres_client.repeat_last_prediction(query_items)
+    query = "SELECT * FROM test_table LIMIT 1"
-    postgres_client.notification_handler.build_and_send_notification.assert_not_called()
+    mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
-    postgres_client.pool.getconn.assert_called_once()
-    postgres_client.pool.putconn.assert_called_once()
 
-    postgres_client.pool.getconn.return_value.cursor.assert_called_once()
+    mock_read_sql_query.return_value = mock_data
-    postgres_client.pool.getconn.return_value.cursor.return_value.execute.assert_called_once_with(
+
-        f"""
+    result = await postgres_activity.load_custom_query(query)
-        INSERT INTO \"{query_items['schema']}\".{query_items['table_name']} (model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, created_at)
+
-        SELECT model_id, prediction, timestamp, response_time, prediction_status, prediction_confidence, NOW()
+    assert isinstance(result, dict)
-        FROM \"{query_items['schema']}\".{query_items['table_name']}
+    assert len(result) == 2
-        WHERE model_id = {query_items['model']}
+    assert "column1" in result
-        ORDER BY timestamp DESC
+    assert "column2" in result
-        LIMIT 1;
+    postgres_activity.logger.info.assert_called()
-        """
-    )
-    postgres_client.pool.getconn.return_value.commit.assert_called_once()
-    postgres_client.pool.getconn.return_value.cursor.return_value.close.assert_called_once()
 
 
 @mark.asyncio
-async def test_repeat_last_prediction_error(postgres_client):
+async def test_load_custom_query_error(postgres_activity):
-    postgres_client.pool.getconn.return_value.cursor.return_value.execute.side_effect = Exception(
+    query = "SELECT * FROM non_existent_table"
-        "Error repeating last prediction")
+    error_msg = "Table not found"
-    query_items = {"schema": "test", "table_name": "test", "model": 1}
+
-    await postgres_client.repeat_last_prediction(query_items)
+    with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
-    postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
+        result = await postgres_activity.load_custom_query(query)
-        notification_id="ERROR_REPEATING_LAST_PREDICTION",
+
-        message="Error repeating last prediction: Error repeating last prediction",
+    assert isinstance(result, dict)
-        block="repeat_last_prediction",
+    assert len(result) == 0
-        level=NotificationLevel.ERROR,
+    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
-        attachment_content=ANY
+    postgres_activity.logger.error.assert_called()
-    )
-    postgres_client.pool.getconn.assert_called_once()
-    postgres_client.pool.putconn.assert_called_once()
 
 
 @mark.asyncio
-@patch("laborious.activities.postgres.DataFrame")
+async def test_repeat_last_prediction_success(postgres_activity):
-async def test_export_data_to_postgres_success(mock_dataframe, postgres_client):
+    query_items = {
-    data = {"schema": "test", "table_name": "test",
+        "schema": "public",
-            "data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
+        "table_name": "predictions",
-    await postgres_client.export_data_to_postgres(data)
+        "model": 1
-    postgres_client.notification_handler.build_and_send_notification.assert_not_called()
+    }
-    postgres_client.pool.getconn.assert_called_once()
-    postgres_client.pool.putconn.assert_called_once()
 
-    mock_dataframe.assert_called_once_with(data["data"])
+    with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
-    mock_dataframe.return_value.to_sql.assert_called_once_with(
+        await postgres_activity.repeat_last_prediction(query_items)
-        data["table_name"],
+
-        postgres_client.pool.getconn.return_value,
+    mock_execute.assert_called_once()
-        schema=data["schema"],
+    postgres_activity.logger.info.assert_called()
-        if_exists="append",
-        index=False
-    )
-    postgres_client.pool.getconn.return_value.commit.assert_called_once()
 
 
 @mark.asyncio
-@patch("laborious.activities.postgres.DataFrame", return_value=MagicMock(
+async def test_repeat_last_prediction_error(postgres_activity):
-    to_sql=MagicMock(side_effect=Exception("Error exporting data to postgres"))
+    query_items = {
-))
+        "schema": "public",
-async def test_export_data_to_postgres_error(mock_dataframe, postgres_client):
+        "table_name": "predictions",
-    data = {"schema": "test", "table_name": "test",
+        "model": 1
-            "data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
+    }
-    await postgres_client.export_data_to_postgres(data)
+    error_msg = "Database error"
-    postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
-        notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
-        message="Error exporting data to postgres: Error exporting data to postgres",
-        block="export_data_to_postgres",
-        level=NotificationLevel.ERROR,
-        attachment_content=ANY
-    )
 
-    postgres_client.pool.getconn.assert_called_once()
+    with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
-    postgres_client.pool.putconn.assert_called_once()
+        await postgres_activity.repeat_last_prediction(query_items)
+
+    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
+    postgres_activity.logger.error.assert_called()
+
+
+@mark.asyncio
+async def test_export_data_to_postgres_success(postgres_activity):
+    input_data = {
+        "schema": "public",
+        "table_name": "test_table",
+        "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
+    }
+
+    with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
+        await postgres_activity.export_data_to_postgres(input_data)
+
+    mock_to_sql.assert_called_once()
+    postgres_activity.logger.debug.assert_called()
+
+
+@mark.asyncio
+async def test_export_data_to_postgres_error(postgres_activity):
+    input_data = {
+        "schema": "public",
+        "table_name": "test_table",
+        "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
+    }
+    error_msg = "Export failed"
+
+    with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
+        await postgres_activity.export_data_to_postgres(input_data)
+
+    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
+    postgres_activity.logger.error.assert_called()
+
+
+@mark.asyncio
+async def test_close(postgres_activity):
+    postgres_activity.close()
+
+    postgres_activity.engine.dispose.assert_called_once()
+
+
+@mark.asyncio
+async def test_del(postgres_activity):
+    postgres_activity.close = MagicMock()
+    postgres_activity.__del__()
+
+    postgres_activity.close.assert_called_once()

tests/laborious/utils/repository/test_opc_repository.py

@@ -181,9 +181,26 @@ def test_validate_connection_lost_time_to_reconect(_mock_datetime, opc_repositor
     assert response == opc_repository.try_connect.return_value
 
 
+def test_validate_connection_failed(opc_repository):
+    opc_repository.client = MagicMock()
+    opc_repository.error_count = 0
+
+    output = opc_repository.validate_connection()
+    assert output is True
+
+
+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.validate_connection.assert_called_once()
+    opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
+
+
 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.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_not_called()

tests/laborious/utils/test_connectors_config.py

@@ -0,0 +1,133 @@
+from os import environ
+from laborious.utils.connectors_config import (build_mlflow_config,
+                                               build_opc_config,
+                                               build_postgres_config)
+
+
+def test_build_mlflow_config_with_env_vars():
+    # Arrange
+    environ['MLFLOW_HOST'] = 'http://test-host'
+    environ['MLFLOW_PORT'] = '8080'
+    environ['MLFLOW_USERNAME'] = 'test-user'
+    environ['MLFLOW_PASSWORD'] = 'test-pass'
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['host'] == 'http://test-host'
+    assert config['port'] == 8080
+    assert config['username'] == 'test-user'
+    assert config['password'] == 'test-pass'
+
+
+def test_build_mlflow_config_with_defaults():
+    # Arrange
+    # Clear any existing env vars
+    environ.pop('MLFLOW_HOST', None)
+    environ.pop('MLFLOW_PORT', None)
+    environ.pop('MLFLOW_USERNAME', None)
+    environ.pop('MLFLOW_PASSWORD', None)
+
+    # Act
+    config = build_mlflow_config()
+
+    # Assert
+    assert config['host'] == 'http://localhost'
+    assert config['port'] == 5080
+    assert config['username'] == 'aignosi'
+    assert config['password'] == 'aignosi'
+
+
+def test_build_opc_config_with_env_vars():
+    # Arrange
+    environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
+
+    # Act
+    config = build_opc_config()
+
+    # Assert
+    assert config['opc']['name'] == 'test-opc'
+    assert config['opc']['url'] == 'opc.tcp://test:4840'
+
+
+def test_build_opc_config_with_individual_env_vars():
+    # Arrange
+    environ.pop('OPC_CONFIG', None)
+    environ['OPC_NAME'] = 'test-name'
+    environ['OPC_URL'] = 'opc.tcp://test:4840'
+    environ['OPC_SERVER_URI'] = 'opc.tcp://test:4840'
+    environ['OPC_RECONNECTION_INTERVAL'] = '300'
+
+    # Act
+    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
+
+
+def test_build_opc_config_with_defaults():
+    # Arrange
+    environ.pop('OPC_CONFIG', None)
+    environ.pop('OPC_NAME', None)
+    environ.pop('OPC_URL', None)
+    environ.pop('OPC_SERVER_URI', None)
+    environ.pop('OPC_RECONNECTION_INTERVAL', None)
+
+    # Act
+    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
+
+
+def test_build_postgres_config_with_env_vars():
+    # Arrange
+    environ['POSTGRES_HOST'] = 'test-host'
+    environ['POSTGRES_PORT'] = '5433'
+    environ['POSTGRES_USER'] = 'test-user'
+    environ['POSTGRES_PASSWORD'] = 'test-pass'
+    environ['POSTGRES_DBNAME'] = 'test-db'
+    environ['POSTGRES_MIN_CONNECTIONS'] = '10'
+    environ['POSTGRES_MAX_CONNECTIONS'] = '30'
+
+    # Act
+    config = build_postgres_config()
+
+    # Assert
+    assert config['host'] == 'test-host'
+    assert config['port'] == 5433
+    assert config['user'] == 'test-user'
+    assert config['password'] == 'test-pass'
+    assert config['dbname'] == 'test-db'
+    assert config['min_connections'] == 10
+    assert config['max_connections'] == 30
+
+
+def test_build_postgres_config_with_defaults():
+    # Arrange
+    environ.pop('POSTGRES_HOST', None)
+    environ.pop('POSTGRES_PORT', None)
+    environ.pop('POSTGRES_USER', None)
+    environ.pop('POSTGRES_PASSWORD', None)
+    environ.pop('POSTGRES_DBNAME', None)
+    environ.pop('POSTGRES_MIN_CONNECTIONS', None)
+    environ.pop('POSTGRES_MAX_CONNECTIONS', None)
+
+    # Act
+    config = build_postgres_config()
+
+    # Assert
+    assert config['host'] == 'localhost'
+    assert config['port'] == 5432
+    assert config['user'] == 'sientia'
+    assert config['password'] == 'sientia'
+    assert config['dbname'] == 'sientia'
+    assert config['min_connections'] == 5
+    assert config['max_connections'] == 20

tests/laborious/utils/test_logger.py

@@ -0,0 +1,37 @@
+import os
+from unittest.mock import patch
+import logging
+import pytest
+from laborious.utils.logger import get_logger
+
+
+@pytest.fixture
+def mock_env_vars():
+    with patch.dict(os.environ, {}, clear=True):
+        yield
+
+
+@pytest.mark.usefixtures("mock_env_vars")
+@patch('laborious.utils.logger.logging.Formatter')
+@patch('laborious.utils.logger.logging.StreamHandler')
+def test_get_logger_defaults(mock_stream_handler, mock_formatter):
+    """Test logger creation with default settings"""
+    # Mock the StreamHandler and Formatter
+
+    logger = get_logger('test_logger')
+
+    # Verify logger settings
+    assert logger.name == 'test_logger'
+    assert logger.level == logging.INFO
+
+    # Verify handler configuration
+    mock_stream_handler.return_value.setLevel.assert_called_once_with('INFO')
+    mock_stream_handler.return_value.setFormatter.assert_called_once()
+
+    # Verify formatter configuration
+    mock_formatter.assert_called_once_with(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+    )
+
+    # Verify handler was added to logger
+    assert len(logger.handlers) == 1

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -55,7 +55,6 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
         call(
             Activities.write_opc_data,
             {
-                'opc_servers': input_data['opc_servers'],
                 'opc_output_config': input_data['opc_output_config'],
                 'data': workflow_mock.execute_local_activity_method.return_value
             },
@@ -116,7 +115,6 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
         call(
             Activities.write_opc_data,
             {
-                'opc_servers': input_data['opc_servers'],
                 'opc_output_config': input_data['opc_output_config'],
                 'data': workflow_mock.execute_local_activity_method.return_value
             },

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -231,7 +231,8 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
         '2024-01-01',  # get_last_timestamp
         ('continue', 0.95, "Input data with bad quality"),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (transform)
+        # mlflow_response_gate (transform)
+        ('continue', 0.95, "Error"),
         # mlflow_content_gate (transform)
         ('continue', 0.95, "Transformed data not passed the content filter"),
     ]
@@ -299,7 +300,8 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
         '2024-01-01',  # get_last_timestamp
         ('continue', 0.95, "Input data with bad quality"),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (transform)
+        # mlflow_response_gate (transform)
+        ('continue', 0.95, "Error"),
         # mlflow_content_gate (transform)
         ('continue', 0.95, "Transformed data not passed the content filter"),
         {'content': 'predicted_data', 'timestamp': '2024-01-01'},  # request_predict
@@ -372,8 +374,14 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, confidence, schema, table_name,
+        data, path_flag, {
-        model, last_timestamp, model_name, model_retention, ""
+            'schema': schema,
+            'table_name': table_name,
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_retention': model_retention
+        }, confidence, last_timestamp, ""
     )
 
     # Assert
@@ -398,8 +406,14 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, confidence, schema, table_name,
+        data, path_flag, {
-        model, last_timestamp, model_name, model_retention, ""
+            'schema': schema,
+            'table_name': table_name,
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_retention': model_retention
+        }, confidence, last_timestamp, ""
     )
 
     # Assert
@@ -433,8 +447,15 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, confidence, schema, table_name,
+        data, path_flag, {
-        model, last_timestamp, model_name, model_retention, 'Prediction Process'
+            'schema': schema,
+            'table_name': table_name,
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_retention': model_retention,
+            'opc_output_config': {'test': 'config'}
+        }, confidence, last_timestamp, 'Prediction Process'
     )
 
     # Assert
@@ -452,7 +473,8 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'model_retention': model_retention,
             'schema': schema,
             'table_name': table_name,
-            'comment': 'Prediction Process'
+            'comment': 'Prediction Process',
+            'opc_output_config': {'test': 'config'}
         }
     )
 
@@ -473,8 +495,15 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, confidence, schema, table_name,
+        data, path_flag, {
-        model, last_timestamp, model_name, model_retention, ""
+            'schema': schema,
+            'table_name': table_name,
+            'model_id': model,
+            'last_timestamp': last_timestamp,
+            'model_name': model_name,
+            'model_retention': model_retention,
+            'opc_output_config': {'test': 'config'}
+        }, confidence, last_timestamp, ""
     )
 
     # Assert

tests/laborious/workflows/test_predictions_batch.py

@@ -55,11 +55,24 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'table_name': input_data['table_name'],
         'model_id': input_data['model_id'],
         'model_name': input_data['model_name'],
-        'input_filters': input_data.get('input_filters', {}),
+        'input_filters': input_data.get('input_filters', {
-        'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {}),
+            'EMPTY_DATA': {
-        'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {}),
+                'POLICY': 'STOP'
+            }
+        }),
+        'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
+            'API_ERROR': {
+                'POLICY': 'STOP'
+            }
+        }),
+        'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
+            'API_ERROR': {
+                'POLICY': 'STOP'
+            }
+        }),
         'model_retention': input_data.get('model_retention', 60),
-        'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
+        'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
+        'opc_output_config': input_data.get('opc_output_config', {})
     }
 
     workflow_mock.execute_child_workflow.assert_has_calls([