gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-10-15 16:50:38 -03:00
parent ac795c7c53
commit f0fb9b854e
14 changed files with 261 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
tests/laborious/activities/test_gates.py
View File

@@ -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

23
tests/laborious/activities/test_mlflow.py
View File

@@ -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 = {

12
tests/laborious/activities/test_storage.py
View File

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

7
tests/laborious/utils/filters/test_conditional_filters.py
View File

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

5
tests/laborious/utils/repository/test_minio_repository.py
View File

@@ -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

57
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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'

24
tests/laborious/utils/repository/test_opc_repository.py
View File

@@ -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

31
tests/laborious/utils/test_connectors_config.py
View File

@@ -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',
}

2
tests/laborious/workflows/test_minimal_retrain.py
View File

@@ -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,