diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 0653161..8cf4afa 100644 --- a/laborious/activities/gates.py +++ b/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) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index d84f527..3e786c7 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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 diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 072274a..be4edb5 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/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: - return False, { - '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) + # ignored because self.validate_connection is called before, so we know self.client is not None + node_obj = self.client.get_node(node) # type: ignore[union-attr] except Exception as e: trace = traceback.format_exc() logger.custom_error(trace, metadata.get('schedule_name', 'N/A')) diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 5497432..02878cf 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/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, diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 1a376e2..8d8822e 100644 --- a/tests/laborious/activities/test_gates.py +++ b/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 diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index cfb66a1..a9f54f8 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/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 = { diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index 87c28f0..76b83e5 100644 --- a/tests/laborious/activities/test_storage.py +++ b/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) diff --git a/tests/laborious/utils/filters/test_conditional_filters.py b/tests/laborious/utils/filters/test_conditional_filters.py index de25b7b..8a35801 100644 --- a/tests/laborious/utils/filters/test_conditional_filters.py +++ b/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( diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py index ec279d2..7883c1c 100644 --- a/tests/laborious/utils/repository/test_minio_repository.py +++ b/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 diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 0e77bdf..a437653 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/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' diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index a0048e8..6ec680d 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/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 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index cf2a6b8..80516ad 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/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', + } diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index c088702..bcbad2d 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/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, diff --git a/validate.sh b/validate.sh index 8728312..aeed301 100755 --- a/validate.sh +++ b/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