SIENTIAPDE-1231

Enhance validation script and refactor code in various modules

- Updated the validation script to include automatic code formatting and linting fixes using Ruff.
- Removed the `clean_tmp_files` method from the Gates class to streamline functionality.
- Simplified conditional checks in the OpcRepository for better clarity and error handling.
- Added model ID to the minimal retrain workflow for improved tracking.
- Introduced new test cases for error handling in MLFlow and storage operations, ensuring robustness in repository interactions.

laborious/activities/gates.py

@@ -3,8 +3,6 @@ from temporalio import activity, workflow
 with workflow.unsafe.imports_passed_through():
     import traceback
     from collections.abc import Callable, Mapping
-    from os import path
-    from shutil import rmtree
     from typing import Any
 
     from pandas import DataFrame
@@ -628,19 +626,3 @@ class Gates(BaseActivity):
         ).observe(response_time)
 
         self.info(f'Metrics written for model {metadata["model_name"]}', metadata)
-
-    @activity.defn(name='clean_tmp_files')
-    async def clean_tmp_files(self, input_data: dict[str, Any]):
-        """
-        Clean temporary files in the tmp directory.
-        """
-        model_name = input_data['model_name']
-        metadata = input_data['metadata']
-        self.info(f'Cleaning tmp files for model {model_name}...', metadata)
-
-        if path.exists(f'tmp/retrain_data/{model_name}'):
-            rmtree(f'tmp/retrain_data/{model_name}')
-        if path.exists(f'tmp/artifacts/{model_name}'):
-            rmtree(f'tmp/artifacts/{model_name}')
-
-        self.info('Tmp files cleaned', metadata)

laborious/utils/repository/model_repository.py

@@ -304,7 +304,7 @@ class MLFlowRepository:
             if model_type == 'predict':
                 model = self.load_predict_model(model_name, flavor)
 
-            elif model_type == 'transform':
+            else:
                 model = self.load_transform_model(model_name, flavor)
 
         return model, artifact_path

laborious/utils/repository/opc_repository.py

@@ -92,14 +92,11 @@ class OpcRepository:
             - Session Timeout: 10,000,000 ms
         """
 
-        if not all([self.cert_path, self.private_key_path]):
+        if self.cert_path is None or self.private_key_path is None:
             raise ValueError(
                 'Certificate and private key paths must be provided for secure connection.'
             )
 
-        if self.cert_path is None or self.private_key_path is None:
-            raise ValueError('Certificate and private key paths cannot be None')
-
         cert = Path(self.cert_path)
         private_key = Path(self.private_key_path)
         server_cert = Path(self.server_cert_path) if self.server_cert_path else None
@@ -316,14 +313,8 @@ class OpcRepository:
         start_time = time.time()
 
         try:
-            if self.client is None:
+            # ignored because self.validate_connection is called before, so we know self.client is not None
-                return False, {
+            node_obj = self.client.get_node(node)  # type: ignore[union-attr]
-                    'notification_id': f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
-                    'message': 'Client is not initialized',
-                    'block': 'opc_repository',
-                    'level': NotificationLevel.ERROR,
-                }
-            node_obj = self.client.get_node(node)
         except Exception as e:
             trace = traceback.format_exc()
             logger.custom_error(trace, metadata.get('schedule_name', 'N/A'))

laborious/workflows/minimal_retrain.py

@@ -116,6 +116,7 @@ class MinimalRetrain:
                 **metadata,
                 'experiment_response': experiment_response,
                 'model_name': model_name,
+                'model_id': input_data['model_id'],
                 'update_report': update_report,
             },
             retry_policy=retry_policy,

tests/laborious/activities/test_gates.py

@@ -117,6 +117,24 @@ async def test_input_gate_with_filter(gates_activity):
     gates_activity.debug.assert_called()
 
 
+@mark.asyncio
+async def test_input_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
+        '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.debug.assert_called()
+
+
 @mark.asyncio
 async def test_mlflow_response_gate_invalid_filter(gates_activity):
     # Arrange
@@ -210,13 +228,38 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
     gates_activity.send_notification.assert_called()
 
 
+@mark.asyncio
+async def test_mlflow_response_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'API_ERROR': {'policy': 'STOP'}},
+        'data': {
+            'success': True,
+            '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.debug.assert_called()
+
+
 @mark.asyncio
 async def test_mlflow_content_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
         'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-        'data': {'value': [1, 2, 3]},
+        'data': {
+            'success': True,
+            'content': {'message': 'success'},
+        },
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -304,6 +347,25 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
     gates_activity.send_notification.assert_called()
 
 
+@mark.asyncio
+async def test_mlflow_content_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'API_ERROR': {'POLICY': 'STOP'}},
+        'data': {'content': {'message': 'success'}},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    }
+
+    # Act
+    result = await gates_activity.mlflow_content_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, '')
+    gates_activity.debug.assert_called()
+
+
 def test_get_prediction_store_policy_invalid_policy(gates_activity):
     # Arrange
     prediction_store_policy = 'INVALID_POLICY'
@@ -506,6 +568,38 @@ async def test_format_default_prediction(gates_activity):
     gates_activity.debug.assert_called()
 
 
+@mark.asyncio
+async def test_format_retrain_report(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'experiment_response': {
+            'success': True,
+            'timestamp': '2023-05-26 11:12:27',
+            'message': 'success',
+        },
+        'update_report': {
+            'version': '1.0.0',
+            'mlflow_run_id': 'test_mlflow_run_id',
+            'mlflow_experiment_id': 'test_mlflow_experiment_id',
+        },
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+    }
+
+    # Act
+    result = await gates_activity.format_retrain_report(input_data)
+
+    # Assert
+    assert result['model_id'] == {0: 'test_model'}
+    assert result['model_name'] == {0: 'test_model'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['status'] == {0: 'success'}
+    assert result['version'] == {0: '1.0.0'}
+    assert result['mlflow_run_id'] == {0: 'test_mlflow_run_id'}
+    assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'}
+
+
 @mark.asyncio
 async def test_get_last_timestamp_with_data(gates_activity):
     # Arrange

tests/laborious/activities/test_mlflow.py

@@ -1,7 +1,7 @@
 from unittest.mock import ANY, MagicMock, call, patch
 
 import numpy as np
-from pytest import fixture, mark
+from pytest import fixture, mark, raises
 from sientia_do.notifications.models import NotificationLevel
 from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
 
@@ -396,6 +396,27 @@ async def test_retrain_model_data_error(mlflow):
     }
 
 
+@mark.asyncio
+async def test_retrain_model_data_error_no_minio_repository(mlflow):
+    mlflow.minio_repository = None
+
+    with raises(ValueError) as e:
+        await mlflow.retrain_model(
+            {
+                **metadata,
+                'object_key': 'test_object_key',
+                'model_name': 'test_model',
+                'model_config': {
+                    'target': 'target',
+                    'transform_flavor': 'sklearn',
+                    'predict_flavor': 'pyfunc',
+                },
+            }
+        )
+
+        assert str(e.value) == 'Minio repository not initialized'
+
+
 @mark.asyncio
 async def test_update_production_model(mlflow):
     input_data = {

tests/laborious/activities/test_storage.py

@@ -1,7 +1,7 @@
 import datetime
 from unittest.mock import ANY, AsyncMock, MagicMock, patch
 
-from pytest import fixture, mark
+from pytest import fixture, mark, raises
 from sientia_do.notifications.models import NotificationLevel
 from sientia_do.temporal.activities.postgres import Postgres
 
@@ -135,6 +135,16 @@ def test___init___done_repository(mock_minio_repository, storage):
     assert storage.minio_repository is not None
 
 
+@mark.asyncio
+async def test_query_to_minio_minio_repository_not_initialized(storage):
+    storage.minio_repository = None
+
+    with raises(ValueError) as e:
+        await storage.query_to_minio({})
+
+        assert str(e.value) == 'Minio repository not initialized'
+
+
 @mark.asyncio
 async def test_query_to_minio_not_data(storage):
     storage.load_custom_query = AsyncMock(return_value=None)

tests/laborious/utils/filters/test_conditional_filters.py

@@ -16,6 +16,13 @@ def test_filter_specific_variables_null_values():
     )
 
 
+def test_filter_specific_variables_null_values_with_empty_data():
+    assert (
+        filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
+        is False
+    )
+
+
 def test_filter_specific_variables_null_values_with_null_values():
     assert (
         filter_specific_variables_null_values(

tests/laborious/utils/repository/test_minio_repository.py

@@ -61,6 +61,11 @@ def minio_repository(mock_boto3, mock_config):
     )
 
 
+def test_close(minio_repository):
+    minio_repository.close()
+    minio_repository.s3_client.close.assert_called_once()
+
+
 def test_ensure_bucket_exists_bucket_exists(minio_repository):
     assert minio_repository.ensure_bucket_exists({}) is True
 

tests/laborious/utils/repository/test_model_repository.py

@@ -173,6 +173,11 @@ def test_get_experiment_none_not_create(mlflow, mlflow_repository):
         assert str(e) == 'Experiment test not found'
 
 
+def test_get_model_params(mlflow, mlflow_repository):
+    output = mlflow_repository.get_model_params('test')
+    assert output == mlflow.get_run.return_value.data.params
+
+
 @patch('laborious.utils.repository.model_repository.path')
 @patch('laborious.utils.repository.model_repository.rmtree')
 @patch('laborious.utils.repository.model_repository.makedirs')
@@ -202,6 +207,35 @@ def test_download_artifacts_success(makedirs, rmtree, path, mlflow_repository):
     assert output == mlflow_repository.client.download_artifacts.return_value
 
 
+@patch('laborious.utils.repository.model_repository.path')
+@patch('laborious.utils.repository.model_repository.rmtree')
+@patch('laborious.utils.repository.model_repository.makedirs')
+def test_download_artifacts_success_path_false(makedirs, rmtree, path, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='test')
+
+    path.exists.return_value = False
+
+    output = mlflow_repository.dowload_artifacts('test', 'path')
+
+    mlflow_repository.get_model_run_id.assert_called_once_with(
+        model_name='test', stage='Production'
+    )
+
+    path.join.assert_called_once_with('./tmp/artifacts/test', 'path')
+
+    path.exists.assert_called_once_with(path.join.return_value)
+
+    rmtree.assert_not_called()
+
+    makedirs.assert_called_once_with('./tmp/artifacts/test', exist_ok=True)
+
+    mlflow_repository.client.download_artifacts.assert_called_once_with(
+        mlflow_repository.get_model_run_id.return_value, 'path', './tmp/artifacts/test'
+    )
+
+    assert output == mlflow_repository.client.download_artifacts.return_value
+
+
 def test_get_experiment_error(mlflow, mlflow_repository):
     mlflow.get_experiment_by_name.return_value = None
 
@@ -383,6 +417,15 @@ valid_cases = [
         },
         ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
     ),
+    (
+        {
+            'value': {
+                datetime(2025, 1, 1, 12, 0, 0, tzinfo=None): 1,
+                datetime(2025, 1, 2, 12, 0, 0, tzinfo=None): 2,
+            }
+        },
+        ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
+    ),
     (
         {
             'value': {
@@ -987,6 +1030,20 @@ def test_retrain_model(mlflow_repository):
     }
 
 
+def test_retrain_model_error(mlflow_repository):
+    data = MagicMock()
+    model_name = 'test'
+    model_config = {'target': 'target', 'transform_flavor': 'sklearn', 'predict_flavor': 'pyfunc'}
+    mlflow_repository.get_model_run_id = MagicMock(side_effect=Exception('error'))
+    output = mlflow_repository.retrain_model(data, model_name, model_config, metadata['metadata'])
+    assert output == {
+        'success': False,
+        'experiment': None,
+        'message': 'Error retraining model test: error',
+        'traceback': ANY,
+    }
+
+
 def test_update_production_model(mlflow_repository):
     experiment = {'run_id': '0', 'experiment_id': '0'}
     model_name = 'test'

tests/laborious/utils/repository/test_opc_repository.py

@@ -86,6 +86,15 @@ async def test_set_security_missing_certificates(opc_repository):
         assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
 
 
+@pytest.mark.asyncio
+async def test_set_security_missing_client(opc_repository):
+    opc_repository.client = None
+    try:
+        await opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Client must be initialized before setting security'
+
+
 @pytest.mark.asyncio
 async def test_connect_with_security(opc_repository, mock_client):
     opc_repository.try_connect = AsyncMock(return_value=(True, {}))
@@ -137,6 +146,21 @@ async def test_try_connect_fail(opc_repository):
     assert error_data['attachment_content'] is not None
 
 
+@pytest.mark.asyncio
+async def test_try_connect_no_client(opc_repository):
+    opc_repository.client = None
+    result = await opc_repository.try_connect()
+    assert result == (
+        False,
+        {
+            'notification_id': f'OPC_CONNECTION_ERROR_{opc_repository.id}',
+            'message': 'Client is not initialized',
+            'block': 'opc_repository',
+            'level': NotificationLevel.ERROR,
+        },
+    )
+
+
 @pytest.mark.asyncio
 async def test_disconnect(opc_repository, mock_client):
     opc_repository.client = mock_client

tests/laborious/utils/test_connectors_config.py

@@ -1,6 +1,7 @@
 from os import environ
 
 from laborious.utils.connectors_config import (
+    build_minio_config,
     build_mlflow_config,
     build_mongodb_config,
     build_opc_config,
@@ -162,3 +163,33 @@ def test_build_mongo_db_config_with_defaults():
         'database_name': 'sientia',
         'ttl_index_seconds': 3600,
     }
+
+
+def test_build_minio_config_with_env_vars():
+    environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
+    environ['MINIO_ACCESS_KEY'] = 'test-key'
+    environ['MINIO_SECRET_KEY'] = 'test-secret'
+    environ['MINIO_REGION_NAME'] = 'test-region'
+    environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
+    assert build_minio_config() == {
+        'endpoint_url': 'http://test-host',
+        'access_key': 'test-key',
+        'secret_key': 'test-secret',
+        'region_name': 'test-region',
+        'default_bucket': 'test-bucket',
+    }
+
+
+def test_build_minio_config_with_defaults():
+    environ.pop('MINIO_ENDPOINT_URL', None)
+    environ.pop('MINIO_ACCESS_KEY', None)
+    environ.pop('MINIO_SECRET_KEY', None)
+    environ.pop('MINIO_REGION_NAME', None)
+    environ.pop('MINIO_DEFAULT_BUCKET', None)
+    assert build_minio_config() == {
+        'endpoint_url': 'http://localhost:9000',
+        'access_key': 'minioadmin',
+        'secret_key': 'minioadmin',
+        'region_name': 'us-east-1',
+        'default_bucket': 'laborious',
+    }

tests/laborious/workflows/test_minimal_retrain.py

@@ -112,6 +112,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
                     **metadata,
                     'experiment_response': {'success': True, 'experiment': 'test_experiment'},
                     'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
                     'update_report': {
                         'success': True,
                         'version': 'test_version',
@@ -267,6 +268,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
                     **metadata,
                     'experiment_response': {'success': False, 'experiment': 'test_experiment'},
                     'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
                     'update_report': {},
                 },
                 retry_policy=ANY,

validate.sh

@@ -47,12 +47,12 @@ run_step() {
 FAILED_STEPS=()
 
 # Step 1: Code Formatting Check (Ruff)
-if ! run_step "1. Code Formatting (Ruff)" "ruff format --check laborious/ tests/"; then
+if ! run_step "1. Code Formatting (Ruff)" "ruff format laborious/ tests/ && ruff format --check laborious/ tests/"; then
     FAILED_STEPS+=("Code Formatting")
 fi
 
 # Step 2: Linting (Ruff)
-if ! run_step "2. Code Linting (Ruff)" "ruff check laborious/ tests/"; then
+if ! run_step "2. Code Linting (Ruff)" "ruff check --fix laborious/ tests/"; then
     FAILED_STEPS+=("Linting")
 fi