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

SIENTIAPDE-1773

Browse Source 
Enhance environment configuration and update dependencies

- Added new environment variables for PluginStore and MLflow configuration in `.env.example`, including `RUNTIME`, `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`, `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `STORE_CACHE_TTL_SECONDS`, `PYPI_SERVER`, `PYPI_USERNAME`, and `PYPI_PASSWORD`.
- Updated `git-requirements-mapping.txt` to reflect changes in repository names.
- Modified `requirements-light.txt` and `requirements.txt` to upgrade `sientia-dataops-library` to version 1.12.0 and `sientia-mlops-library` to version 0.8.1.
- Updated `values.yaml` to include new environment variables for worker runtime and PluginStore configuration.
- Refactored E2E tests to utilize new MLflow repository stubs and PluginStore mocks for improved testing accuracy.
This commit is contained in:
vitor-aignosi
2026-05-05 16:52:51 -03:00
parent 473bd0b03f
commit 1ce8b9d3a7
23 changed files with 1280 additions and 3970 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/laborious/activities/test_activities.py
View File

@@ -48,7 +48,8 @@ def test___init__(
'secure': False,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
mlflow_repository = MagicMock()
plugin_store = MagicMock()
opc_config = {
'bootstrap_servers': 'localhost:9092',
@@ -67,12 +68,13 @@ def test___init__(
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
plugin_store=plugin_store,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler,
mlflow_repository=mlflow_repository,
)
assert isinstance(activities, Activities)
@@ -101,10 +103,8 @@ def test___init__(
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'],
mlflow_repository=mlflow_repository,
plugin_store=plugin_store,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
@@ -193,7 +193,8 @@ async def test_shutdown(
'secure': False,
}
mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
mlflow_repository = MagicMock()
plugin_store = MagicMock()
opc_config = {
'bootstrap_servers': 'localhost:9092',
@@ -212,12 +213,13 @@ async def test_shutdown(
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
plugin_store=plugin_store,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler,
mlflow_repository=mlflow_repository,
)
await activities.shutdown()

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

@@ -1,9 +1,10 @@
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import numpy as np
import pandas as pd
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
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
@@ -16,12 +17,14 @@ def _passthrough_from_dict():
yield
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository, mock_mlflow_repository):
def test___init__(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = AsyncMock()
mlflow_repo = MagicMock()
plugin_store = MagicMock()
minio_repo = mock_minio_repository(
endpoint='localhost:9000',
access_key='minio',
@@ -32,24 +35,16 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
bucket='test',
)
mlflow = MLFlow(
mlflow_host='http://localhost',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
mlflow_repository=mlflow_repo,
plugin_store=plugin_store,
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
assert mlflow.mlflow_host == 'http://localhost'
assert mlflow.mlflow_port == 5000
assert mlflow.mlflow_username == 'admin'
assert mlflow.mlflow_password == 'admin'
mock_mlflow_repository.assert_called_once_with(
'http://localhost:5000', 'admin', 'admin', ANY, ANY, ANY
)
assert mlflow.mlflow_repository is mlflow_repo
assert mlflow.plugin_store is plugin_store
mock_minio_repository.assert_called_once_with(
endpoint='localhost:9000',
@@ -63,12 +58,14 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
@fixture
@patch('laborious.activities.mlflow.MLFlowRepository')
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository, mock_mlflow_repository):
def mlflow(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = AsyncMock()
mlflow_repo = MagicMock()
plugin_store = MagicMock()
minio_repo = mock_minio_repository(
endpoint='localhost:9000',
access_key='minio',
@@ -79,17 +76,14 @@ def mlflow(mock_minio_repository, mock_mlflow_repository):
bucket='test',
)
mlflow = MLFlow(
mlflow_host='http://localhost:5000',
mlflow_port=5000,
mlflow_username='admin',
mlflow_password='admin',
mlflow_repository=mlflow_repo,
plugin_store=plugin_store,
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
mlflow.model_monitoring_repository = AsyncMock()
mlflow.minio_repository = AsyncMock()
mlflow.send_notification = MagicMock()
@@ -120,9 +114,32 @@ metadata = {
new_callable=AsyncMock,
)
async def test_request_transform_success(mock_from_dataframe, mlflow):
data_mock = MagicMock()
ts = pd.Timestamp('2020-01-01', tz='UTC')
raw = pd.DataFrame(
{
'variable': ['v1', 'v1'],
'timestamp': [ts, ts],
'value': [1.0, 2.0],
'created_at': [ts, ts],
}
)
pivoted = raw.sort_values('created_at', ascending=False).drop_duplicates(
subset=['variable', 'timestamp'], keep='first'
)
pivoted = pivoted.pivot(index='timestamp', columns='variable', values='value')
pivoted = pivoted.fillna(np.nan)
pivoted.columns.name = None
pivoted.index.name = None
pivoted['timestamp'] = pivoted.index
out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
wrapper = MagicMock()
wrapper.transform.return_value = (out_df, {'meta': True})
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=data_mock)
payload.retrieve = AsyncMock(return_value=raw)
input_data = {
**metadata,
@@ -131,17 +148,13 @@ async def test_request_transform_success(mock_from_dataframe, mlflow):
'model_config': {},
}
transform_response = {'success': True, 'content': MagicMock()}
mlflow.model_monitoring_repository.transform.return_value = transform_response
data_mock.sort_values.return_value = data_mock
data_mock.drop_duplicates.return_value = data_mock
data_mock.pivot.return_value = data_mock
response_data = await mlflow.request_transform(input_data)
mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', data_mock, {}, metadata['metadata']
mlflow.mlflow_repository.get_cached_model.assert_called_once_with(
model_name='test_model',
alias='production',
retention_minutes=0,
metadata=metadata['metadata'],
)
mock_from_dataframe.assert_called_once()
assert response_data == mock_from_dataframe.return_value
@@ -153,6 +166,8 @@ async def test_request_transform_success(mock_from_dataframe, mlflow):
new_callable=AsyncMock,
)
async def test_request_transform_failure(mock_from_dataframe, mlflow):
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('boom')
data_mock = MagicMock()
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=data_mock)
@@ -164,26 +179,22 @@ async def test_request_transform_failure(mock_from_dataframe, mlflow):
'model_config': {},
}
transform_response = {'success': False, 'message': 'Transform failed'}
mlflow.model_monitoring_repository.transform.return_value = transform_response
data_mock.sort_values.return_value = data_mock
data_mock.drop_duplicates.return_value = data_mock
data_mock.pivot.return_value = data_mock
response_data = await mlflow.request_transform(input_data)
await mlflow.request_transform(input_data)
mock_from_dataframe.assert_called_once_with(
dataframe=None,
minio_repo=mlflow.minio_repository,
model_name='test_model',
operation='transform',
status=transform_response,
status={'success': False, 'content': ANY},
workflow_metadata=metadata['metadata'],
last_timestamp=payload.last_timestamp,
logger=mlflow.logger,
)
assert response_data == mock_from_dataframe.return_value
@mark.asyncio
@@ -193,7 +204,13 @@ async def test_request_transform_failure(mock_from_dataframe, mlflow):
)
@patch('laborious.activities.mlflow.to_datetime')
async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
wrapper = MagicMock()
pred_df = MagicMock()
wrapper.predict.return_value = (pred_df, {})
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
data_mock = MagicMock()
data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=data_mock)
@@ -204,20 +221,14 @@ async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
'model_config': {},
}
predict_response = {'success': True, 'content': MagicMock()}
mlflow.model_monitoring_repository.predict.return_value = predict_response
pred_df.columns = MagicMock()
pred_df.__setitem__ = MagicMock()
response_data = await mlflow.request_predict(input_data)
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
mock_to_datetime.assert_called_once_with(
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', data_mock, {}, metadata['metadata']
)
mock_to_datetime.assert_called()
mlflow.mlflow_repository.get_cached_model.assert_called_once()
mock_from_dataframe.assert_called_once()
assert response_data == mock_from_dataframe.return_value
@@ -229,6 +240,8 @@ async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
)
@patch('laborious.activities.mlflow.to_datetime')
async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('predict boom')
data_mock = MagicMock()
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=data_mock)
@@ -240,34 +253,51 @@ async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, ml
'model_config': {},
}
predict_response = {'success': False, 'message': 'Predict failed'}
mlflow.model_monitoring_repository.predict.return_value = predict_response
response_data = await mlflow.request_predict(input_data)
await mlflow.request_predict(input_data)
mock_from_dataframe.assert_called_once_with(
dataframe=None,
minio_repo=mlflow.minio_repository,
model_name='test_model',
operation='predict',
status=predict_response,
status={'success': False, 'content': ANY},
workflow_metadata=metadata['metadata'],
last_timestamp=payload.last_timestamp,
logger=mlflow.logger,
)
assert response_data == mock_from_dataframe.return_value
@mark.asyncio
@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
}
async def test_retrain_model_success_data_success_retrain(
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
):
mock_mkdtemp.return_value = '/tmp/x'
mv_alias = MagicMock()
mv_alias.run_id = 'source-run'
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
wrapper = MagicMock()
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
mock_cm = MagicMock()
mock_cm.__enter__.return_value = MagicMock(run_id='new-run', experiment_id='exp-1')
mock_cm.__exit__.return_value = False
mlflow.mlflow_repository.start_run.return_value = mock_cm
ts = pd.Timestamp('2020-01-01', tz='UTC')
raw_data = pd.DataFrame(
{
'variable': ['target', 'f1'],
'timestamp': [ts, ts],
'value': [1.0, 2.0],
}
)
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=raw_data)
@@ -278,177 +308,93 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_not_called()
raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
raw_data = raw_data.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
assert response == {
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
'timestamp': timestamp,
}
wrapper.retrain.assert_called_once()
wrapper.store_model.assert_called_once_with(name='test_model')
assert response['success'] is True
assert response['experiment']['run_id'] == 'new-run'
@mark.asyncio
@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': True,
'experiment': 'test_experiment',
'message': 'Model retrained successfully.',
}
async def test_retrain_model_success_with_payload_data(
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
):
mock_mkdtemp.return_value = '/tmp/x'
mv_alias = MagicMock(run_id='src')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
wrapper = MagicMock()
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
mock_cm = MagicMock()
mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
mock_cm.__exit__.return_value = False
mlflow.mlflow_repository.start_run.return_value = mock_cm
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
raw_data.__getitem__.return_value.max.return_value = 'ts'
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=raw_data)
pivoted = MagicMock()
raw_data.sort_values.return_value = raw_data
raw_data.drop_duplicates.return_value = raw_data
raw_data.pivot.return_value = pivoted
pivoted.fillna = MagicMock()
pivoted.columns.name = None
pivoted.index = MagicMock()
pivoted.__setitem__ = MagicMock()
response = await mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
'model_config': {'target': 'target'},
}
)
assert response['success'] is True
mlflow.minio_repository.download_file.assert_not_called()
@mark.asyncio
@patch('laborious.activities.mlflow.to_datetime')
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
mlflow.model_monitoring_repository.retrain_model.return_value = {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
}
mv_alias = MagicMock(run_id='src')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('retrain failed')
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
raw_data.__getitem__.return_value.max.return_value = 'tsmax'
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=raw_data)
pivoted = MagicMock()
raw_data.sort_values.return_value = raw_data
raw_data.drop_duplicates.return_value = raw_data
raw_data.pivot.return_value = pivoted
pivoted.fillna = MagicMock()
pivoted.columns.name = None
pivoted.index = MagicMock()
pivoted.__setitem__ = MagicMock()
response = await mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
'model_config': {'target': 'target'},
}
)
timestamp = raw_data.__getitem__.return_value.max.return_value
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
subset=['variable', 'timestamp'], keep='first'
)
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
raw_data.drop.assert_has_calls(
[
call(columns=['model_id'], inplace=True, errors='ignore'),
call(columns=['created_at'], inplace=True, errors='ignore'),
]
)
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
raw_data = raw_data.pivot.return_value
raw_data.__setitem__.assert_has_calls(
[
call('timestamp', raw_data.index),
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
call('timestamp', mock_to_datetime.return_value),
]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
)
mock_to_datetime.assert_has_calls(
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
)
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
data=raw_data,
model_name='test_model',
model_config={
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
metadata=metadata['metadata'],
)
mlflow.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='RETRAIN_MODEL_ERROR',
message='Error retraining model test_model: Model retrained failed.',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
assert response == {
'success': False,
'traceback': 'test_traceback',
'message': 'Model retrained failed.',
'timestamp': timestamp,
}
assert response['success'] is False
mlflow.send_notification_async.assert_called_once()
assert 'retrain failed' in response['message']
@mark.asyncio
@@ -459,18 +405,40 @@ async def test_retrain_model_data_error(mlflow):
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
assert response == {
'success': False,
'message': "Error loading retrain data: 'data'",
'traceback': ANY,
'timestamp': ANY,
}
assert response['success'] is False
assert 'data' in response['message'].lower() or 'loading' in response['message'].lower()
@mark.asyncio
async def test_retrain_model_missing_target(mlflow):
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
raw_data.__getitem__.return_value.max.return_value = 'ts'
payload = AsyncMock()
payload.retrieve = AsyncMock(return_value=raw_data)
pivoted = MagicMock()
raw_data.drop_duplicates.return_value = raw_data
raw_data.pivot.return_value = pivoted
pivoted.fillna = MagicMock()
pivoted.columns.name = None
pivoted.index = MagicMock()
pivoted.__setitem__ = MagicMock()
response = await mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
)
assert response['success'] is False
assert 'target' in response['message']
@mark.asyncio
@@ -485,96 +453,94 @@ async def test_retrain_model_data_error_no_minio_repository(mlflow):
'model_name': 'test_model',
'model_config': {
'target': 'target',
'transform_flavor': 'sklearn',
'predict_flavor': 'pyfunc',
},
}
)
assert str(e.value) == 'Minio repository not initialized'
assert str(e.value) == 'Minio repository not initialized'
@mark.asyncio
async def test_update_production_model(mlflow):
mlflow.mlflow_repository._client.search_model_versions.return_value = [
MagicMock(version='3', run_id='run-x'),
MagicMock(version='2', run_id='run-x'),
]
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
'timestamp': 2,
'status': 'success',
}
response = await mlflow.update_production_model(input_data)
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
experiment='test', model_name='test_model', metadata=metadata['metadata']
mlflow.mlflow_repository.promote_to_alias.assert_called_once_with(
model_name='test_model',
version='3',
alias='production',
metadata=metadata['metadata'],
)
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
assert response['model_name'] == 'test_model'
assert response['version'] == '3'
@mark.asyncio
async def test_update_production_model_error(mlflow):
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
'Error updating production model'
)
mlflow.mlflow_repository._client.search_model_versions.return_value = []
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': 'test',
'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
'timestamp': 2,
'status': 'success',
}
try:
await mlflow.update_production_model(input_data)
except Exception as e:
assert str(e) == 'Error updating production model'
except Exception:
mlflow.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message='Error updating production model test_model: Error updating production model',
message=ANY,
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('No exception raised')
raise AssertionError('Expected exception')
@mark.asyncio
@patch('laborious.activities.mlflow.to_datetime')
async def test_get_reference_data_success(mock_to_datetime, mlflow):
# Arrange
input_data = {
**metadata,
'model_name': 'test_model',
}
# Mock reference data DataFrame
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mock_reference_data = MagicMock()
mock_reference_data.__getitem__.return_value = MagicMock()
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
mock_reference_data.to_dict.return_value = [
{'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
{'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
]
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
# Act
result = await mlflow.get_reference_data(input_data)
# Assert
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
model_name='test_model',
artifact_path='evaluation_data.csv',
metadata=metadata['metadata'],
)
mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
with patch('laborious.activities.mlflow.rmtree'):
with patch('laborious.activities.mlflow.Path') as mp:
mp.return_value.rglob.return_value = [MagicMock()]
result = await mlflow.get_reference_data(input_data)
mock_reference_data.to_dict.assert_called_once_with(orient='records')
assert result == mock_reference_data.to_dict.return_value
@@ -582,48 +548,30 @@ async def test_get_reference_data_success(mock_to_datetime, mlflow):
@mark.asyncio
async def test_get_reference_data_not_found(mlflow):
# Arrange
input_data = {
**metadata,
'model_name': 'test_model',
}
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
mlflow.mlflow_repository._client.get_model_version_by_alias.side_effect = Exception('missing')
# Act
result = await mlflow.get_reference_data(input_data)
# Assert
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
model_name='test_model',
artifact_path='evaluation_data.csv',
metadata=metadata['metadata'],
)
mlflow.warning.assert_called_once_with(
'Reference data not found for model test_model', metadata['metadata']
)
mlflow.warning.assert_called()
assert result is None
@mark.asyncio
async def test_get_reference_data_exception(mlflow):
# Arrange
input_data = {
**metadata,
'model_name': 'test_model',
}
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
'Error loading artifact'
)
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
# Act & Assert
with raises(Exception) as e:
await mlflow.get_reference_data(input_data)
result = await mlflow.get_reference_data(input_data)
assert str(e.value) == 'Error loading artifact'
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
model_name='test_model',
artifact_path='evaluation_data.csv',
metadata=metadata['metadata'],
)
assert result is None

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

File diff suppressed because it is too large Load Diff

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

@@ -4,6 +4,7 @@ from laborious.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_opc_config,
build_plugin_store_config,
)
@@ -18,12 +19,22 @@ def test_build_mlflow_config_with_env_vars():
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://test-host'
assert config['port'] == 8080
assert config['url'] == 'http://test-host:8080'
assert config['username'] == 'test-user'
assert config['password'] == 'test-pass'
def test_build_mlflow_config_host_already_has_port():
environ['MLFLOW_HOST'] = 'http://tracker.example.com:443'
environ['MLFLOW_PORT'] = '8080'
environ['MLFLOW_USERNAME'] = 'u'
environ['MLFLOW_PASSWORD'] = 'p'
config = build_mlflow_config()
assert config['url'] == 'http://tracker.example.com:443'
def test_build_mlflow_config_with_defaults():
# Arrange
# Clear any existing env vars
@@ -36,12 +47,30 @@ def test_build_mlflow_config_with_defaults():
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://localhost'
assert config['port'] == 5080
assert config['url'] == 'http://localhost:5080'
assert config['username'] == 'aignosi'
assert config['password'] == 'aignosi'
def test_build_plugin_store_config_defaults():
environ.pop('STORE_BASE_URL', None)
environ.pop('STORE_OWNER', None)
environ.pop('STORE_REPO', None)
environ.pop('STORE_BRANCH', None)
environ.pop('STORE_USERNAME', None)
environ.pop('STORE_PASSWORD', None)
environ.pop('STORE_CACHE_TTL_SECONDS', None)
environ.pop('PYPI_SERVER', None)
environ.pop('PYPI_USERNAME', None)
environ.pop('PYPI_PASSWORD', None)
cfg = build_plugin_store_config()
assert cfg['base_url'] == 'http://localhost:3000'
assert cfg['owner'] == 'sientia'
assert cfg['repo'] == 'model-library-store'
assert cfg['pypi_index_url'] == 'http://localhost:5000'
def test_build_opc_config_with_env_vars():
# Arrange
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'

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

@@ -34,8 +34,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
'retention_minutes': 0,
},
}
@@ -165,8 +164,7 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
'retention_minutes': 0,
},
}
@@ -227,8 +225,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'transform_flavor': 'test_transform_flavor',
'predict_flavor': 'test_predict_flavor',
'retention_minutes': 0,
},
}