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

Code import - branch release/SIENTIAPDE-1646

Browse Source
This commit is contained in:
gitea_admin
2026-08-05 13:53:38 +00:00
commit a8e89535ee
106 changed files with 24108 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/laborious/activities/__init__.py Normal file
View File

299
tests/laborious/activities/test_activities.py Normal file
View File

@@ -0,0 +1,299 @@
from unittest.mock import ANY, MagicMock, patch
from laborious.activities.activities import Activities
from laborious.activities.api import API
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
@patch('laborious.activities.activities.Storage.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test___init__(
mock_metrics_controller,
mock_minio_repository,
mock_api_init,
mock_model_metrics_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
mock_storage_init,
):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10,
}
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'default_bucket': 'test',
'retention_hours': 24,
'secure': False,
}
mlflow_repository = MagicMock()
plugin_store = MagicMock()
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_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)
assert isinstance(activities, Storage)
assert isinstance(activities, MLFlow)
assert isinstance(activities, OPC)
assert isinstance(activities, Gates)
assert isinstance(activities, ModelMetrics)
assert isinstance(activities, API)
mock_storage_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'],
retention_hours=minio_config['retention_hours'],
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_repository=mlflow_repository,
plugin_store=plugin_store,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_opc_init.assert_called_once_with(
ANY,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_gates_init.assert_called_once_with(
ANY,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_model_metrics_init.assert_called_once_with(
ANY,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_api_init.assert_called_once_with(
ANY,
base_url=pi_web_api_config['base_url'],
auth_type=pi_web_api_config['auth_type'],
auth_token=pi_web_api_config['auth_token'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_minio_repository.assert_called_once_with(
endpoint=minio_config['endpoint_url'],
access_key=minio_config['access_key'],
secret_key=minio_config['secret_key'],
bucket=minio_config['default_bucket'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
secure=minio_config['secure'],
)
@patch('laborious.activities.activities.Storage')
@patch('laborious.activities.activities.MLFlow')
@patch('laborious.activities.activities.OPC')
@patch('laborious.activities.activities.Gates')
@patch('laborious.activities.activities.ModelMetrics')
@patch('laborious.activities.activities.API')
@patch('laborious.activities.activities.MinioRepository')
def test_shutdown(
_mock_minio_repository,
mock_api_init,
mock_model_metrics_init,
mock_gates_init,
mock_opc_init,
mock_mlflow_init,
mock_storage_init,
):
mock_opc_init.close = MagicMock()
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10,
}
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'default_bucket': 'test',
'retention_hours': 24,
'secure': False,
}
mlflow_repository = MagicMock()
plugin_store = MagicMock()
opc_config = {
'bootstrap_servers': 'localhost:9092',
'polling_time': 1000,
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
notification_handler = MagicMock()
activities = Activities(
postgres_config=postgres_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,
)
activities.shutdown()
mock_opc_init.close.assert_called_once()
mock_storage_init.close.assert_called_once()
mock_mlflow_init.close.assert_called_once()
mock_gates_init.close.assert_called_once()
mock_model_metrics_init.close.assert_called_once()
mock_api_init.close.assert_called_once()
@patch('laborious.activities.activities.SientiaMLflowRepository')
@patch('laborious.activities.activities.build_mlflow_config')
@patch('laborious.activities.activities.Storage.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
@patch('laborious.activities.activities.ModelMetrics.__init__')
@patch('laborious.activities.activities.API.__init__')
@patch('laborious.activities.activities.MinioRepository')
@patch('laborious.activities.activities.MetricsController')
def test___init___builds_mlflow_repository_when_not_provided(
mock_metrics_controller,
mock_minio_repository,
_mock_api_init,
_mock_model_metrics_init,
_mock_gates_init,
_mock_opc_init,
_mock_mlflow_init,
_mock_storage_init,
mock_build_mlflow_config,
mock_mlflow_repository_cls,
):
postgres_config = {
'host': 'localhost',
'port': 5432,
'user': 'postgres',
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10,
}
minio_config = {
'endpoint_url': 'localhost:9000',
'access_key': 'minio',
'secret_key': 'minio123',
'default_bucket': 'test',
'retention_hours': 24,
'secure': False,
}
opc_config = {'bootstrap_servers': 'localhost:9092', 'polling_time': 1000, 'group_id': 'test'}
pi_web_api_config = {'base_url': 'https://pi', 'auth_type': 'bearer', 'auth_token': 'token'}
logger = MagicMock()
notification_handler = MagicMock()
plugin_store = MagicMock()
mock_build_mlflow_config.return_value = {
'url': 'http://mlflow:80',
'username': 'u',
'password': 'p',
}
Activities(
postgres_config=postgres_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,
)
mock_build_mlflow_config.assert_called_once()
mock_mlflow_repository_cls.assert_called_once_with(
host='http://mlflow:80',
username='u',
password='p',
logger=logger,
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)

503
tests/laborious/activities/test_api.py Normal file
View File

@@ -0,0 +1,503 @@
from unittest.mock import ANY, MagicMock, call, patch
from pytest import fixture
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def _create_mock_dataframe(to_dict_return=None):
"""Helper function to create a mocked DataFrame for testing."""
mock_df = MagicMock()
mock_head = MagicMock()
def get_column_values(key):
if key == 'prediction':
return MagicMock(values=[0.75])
elif key == 'prediction_confidence':
return MagicMock(values=[0.95])
else:
return MagicMock(values=['2024-01-01T00:00:00+00:00'])
mock_head.__getitem__.side_effect = get_column_values
mock_df.head.return_value = mock_head
if to_dict_return is None:
to_dict_return = {
'prediction': [0.75],
'prediction_confidence': [0.95],
'timestamp': ['2024-01-01T00:00:00+00:00'],
}
mock_df.to_dict.return_value = to_dict_return
return mock_df
@fixture
def base_input_data():
"""Base input data for PI Web API tests."""
return {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95],
'timestamp': ['2024-01-01T00:00:00+00:00'],
},
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com/piwebapi',
'prediction_tags': {'tag1': 'web_id_1'},
'confidence_tags': {'tag2': 'web_id_2'},
},
}
@patch('laborious.activities.api.PIWebAPIClient')
def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client):
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
api_instance = API(
base_url='https://test-pi-server.com',
auth_type='bearer',
auth_token='test_token',
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
with patch.object(
SientiaMonitoring,
'get_core_labels',
return_value={
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'operation_type': 'write_pi_web_api_data',
},
):
labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata'])
assert labels['operation_type'] == 'write_pi_web_api_data'
assert labels == {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'operation_type': 'write_pi_web_api_data',
}
@patch('laborious.activities.api.PIWebAPIClient')
def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
api_instance = API(
base_url='https://test-pi-server.com',
auth_type='bearer',
auth_token='test_token',
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
with patch.object(
SientiaMonitoring,
'get_core_labels',
return_value={
'pod_id': 'test_pod',
'runtime': 'k8s',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'operation_type': 'write',
},
):
labels = api_instance.get_pi_web_api_core_labels(
metadata=metadata['metadata'], operation_type='write'
)
assert labels['operation_type'] == 'write'
assert labels['runtime'] == 'k8s'
def test__init__():
api = API(
base_url='https://test-pi-server.com',
auth_type='bearer',
auth_token='test_token',
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
assert api.pi_web_api_client is not None
@fixture
@patch('laborious.activities.api.PIWebAPIClient')
def api(mock_pi_web_api_client):
mock_client = MagicMock()
mock_client.write_value = MagicMock()
mock_client.close = MagicMock()
mock_client.base_url = 'https://test-pi-server.com'
mock_pi_web_api_client.return_value = mock_client
api_instance = API(
base_url='https://test-pi-server.com',
auth_type='bearer',
auth_token='test_token',
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
api_instance.send_notification = MagicMock()
api_instance.info = MagicMock()
api_instance.error = MagicMock()
api_instance.emit_metric_sync = MagicMock()
api_instance.get_core_labels = MagicMock(
return_value={
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
)
return api_instance
@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
input_data = {
**base_input_data,
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com/piwebapi',
'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
},
}
mock_dataframe.return_value = _create_mock_dataframe()
# Mock successful responses
api.pi_web_api_client.write_value.side_effect = [
[{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
[{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
]
result = api.write_pi_web_api_data(input_data)
api.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=['web_id_1', 'web_id_2'],
value={
'Timestamp': '2024-01-01T00:00:00+00:00',
'Value': 0.75,
},
metadata=metadata['metadata'],
),
call(
web_ids=['web_id_3', 'web_id_4'],
value={
'Timestamp': '2024-01-01T00:00:00+00:00',
'Value': 0.95,
},
metadata=metadata['metadata'],
),
]
)
assert result == {
'prediction': [0.75],
'prediction_confidence': [0.95],
'timestamp': ['2024-01-01T00:00:00+00:00'],
}
@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
mock_dataframe.return_value = _create_mock_dataframe(
{
'prediction': [0.75],
'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
'timestamp': ['2024-01-01T00:00:00+00:00'],
}
)
api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')
result = api.write_pi_web_api_data(base_input_data)
api.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert api.pi_web_api_client.write_value.call_count == 1
@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
mock_dataframe.return_value = _create_mock_dataframe()
# First call succeeds, second fails
api.pi_web_api_client.write_value.side_effect = [
[{'WebId': 'web_id_1', 'Errors': []}],
Exception('Confidence write failed'),
]
result = api.write_pi_web_api_data(base_input_data)
api.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
assert result == {
'prediction': [0.75],
'prediction_confidence': [0.95],
'timestamp': ['2024-01-01T00:00:00+00:00'],
}
assert api.pi_web_api_client.write_value.call_count == 2
@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
input_data = {
**base_input_data,
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com/piwebapi',
'prediction_tags': {},
'confidence_tags': {},
},
}
mock_dataframe.return_value = _create_mock_dataframe()
# Mock empty responses
api.pi_web_api_client.write_value.side_effect = [
[],
[],
]
result = api.write_pi_web_api_data(input_data)
api.pi_web_api_client.write_value.assert_has_calls(
[
call(
web_ids=[],
value={
'Timestamp': '2024-01-01T00:00:00+00:00',
'Value': 0.75,
},
metadata=metadata['metadata'],
),
call(
web_ids=[],
value={
'Timestamp': '2024-01-01T00:00:00+00:00',
'Value': 0.95,
},
metadata=metadata['metadata'],
),
]
)
assert result == {
'prediction': [0.75],
'prediction_confidence': [0.95],
'timestamp': ['2024-01-01T00:00:00+00:00'],
}
@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_updates_confidence_and_comments(
mock_dataframe, api, base_input_data
):
mock_dataframe.return_value = _create_mock_dataframe()
api.pi_web_api_client.write_value.side_effect = [
[{'WebId': 'web_id_1', 'Errors': []}],
[{'WebId': 'web_id_2', 'Errors': []}],
]
with patch.object(
api,
'process_pi_web_api_response',
new=MagicMock(side_effect=[(0.33, 'PI warning'), (0, '')]),
) as process_mock:
result = api.write_pi_web_api_data(base_input_data)
assert process_mock.call_count == 2
assert result is not None
@patch('laborious.activities.api.SientiaMonitoring.shutdown')
def test_close(mock_shutdown, api):
api.close()
api.pi_web_api_client.close.assert_called_once()
mock_shutdown.assert_called_once_with(api)
def test_process_pi_web_api_response_success(api):
"""Test successful processing of PI Web API response with all tags written."""
response_data = [
{'WebId': 'web_id_1', 'Errors': []},
{'WebId': 'web_id_2', 'Errors': []},
]
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
confidence, message = api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
metadata=metadata['metadata'],
)
assert confidence == 0
assert message == ''
assert api.emit_metric_sync.call_count == 2
# Verify that emit_metric_sync was called with correct tags structure
call_args_list = api.emit_metric_sync.call_args_list
assert len(call_args_list) == 2
# Check that all calls include core_labels and tag_name
for call_args in call_args_list:
assert 'tag_name' in call_args.kwargs['tags']
assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']
def test_process_pi_web_api_response_with_errors(api):
"""Test processing response with errors in some tags."""
response_data = [
{'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
{'WebId': 'web_id_2', 'Errors': []},
]
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
confidence, message = api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
metadata=metadata['metadata'],
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert (
message
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
)
assert api.emit_metric_sync.call_count == 2
def test_process_pi_web_api_response_missing_tags(api):
"""Test processing response when number of written tags doesn't match expected."""
response_data = [
{'WebId': 'web_id_1', 'Errors': []},
]
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
confidence, message = api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
metadata=metadata['metadata'],
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert (
message
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
)
api.send_notification.assert_called_once()
call_args = api.send_notification.call_args
assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
assert call_args.kwargs['level'] == NotificationLevel.ERROR
def test_process_pi_web_api_response_missing_webid(api):
"""Test processing response when WebId is missing in response item."""
response_data = [
{'Errors': []},
{'WebId': 'web_id_2', 'Errors': []},
]
tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
confidence, message = api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
metadata=metadata['metadata'],
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert (
message
== "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
)
api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])
def test_process_pi_web_api_response_missing_tag_name(api):
"""Test processing response when tag name is not found for WebId."""
response_data = [
{'WebId': 'unknown_web_id', 'Errors': []},
]
tags = {'tag1': 'web_id_1'}
core_labels = {
'pod_id': 'test_pod',
'runtime': 'local',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
confidence, message = api.process_pi_web_api_response(
response_data=response_data,
tags=tags,
core_labels=core_labels,
metadata=metadata['metadata'],
)
assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
assert (
message
== "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
)
api.error.assert_any_call(
'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
)

1006
tests/laborious/activities/test_gates.py Normal file
View File

File diff suppressed because it is too large Load Diff

856
tests/laborious/activities/test_mlflow.py Normal file
View File

@@ -0,0 +1,856 @@
from datetime import datetime
from unittest.mock import ANY, MagicMock, patch
import numpy as np
import pandas as pd
from pytest import fixture, raises
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow
@fixture(autouse=True)
def _passthrough_from_dict():
with patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
):
yield
@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = MagicMock()
mlflow_repo = MagicMock()
plugin_store = MagicMock()
minio_repo = mock_minio_repository(
endpoint='localhost:9000',
access_key='minio',
secret_key='minio123',
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
bucket='test',
)
mlflow = MLFlow(
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_repository is mlflow_repo
assert mlflow.plugin_store is plugin_store
mock_minio_repository.assert_called_once_with(
endpoint='localhost:9000',
access_key='minio',
secret_key='minio123',
logger=ANY,
notification_handler=ANY,
metrics_controller=ANY,
bucket='test',
)
@fixture
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = MagicMock()
mlflow_repo = MagicMock()
plugin_store = MagicMock()
minio_repo = mock_minio_repository(
endpoint='localhost:9000',
access_key='minio',
secret_key='minio123',
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
bucket='test',
)
mlflow = MLFlow(
mlflow_repository=mlflow_repo,
plugin_store=plugin_store,
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
mlflow.minio_repository = MagicMock()
mlflow.send_notification = MagicMock()
mlflow.emit_metric = MagicMock()
mlflow.error = MagicMock()
mlflow.debug = MagicMock()
mlflow.info = MagicMock()
mlflow.warning = MagicMock()
mlflow.critical = MagicMock()
return mlflow
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test_detect_and_parse_datetime_index_empty(mlflow):
df = pd.DataFrame()
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
assert out.empty
def test_detect_and_parse_datetime_index_mixed_types_error(mlflow):
idx = pd.Index([pd.Timestamp('2020-01-01', tz='UTC'), 'x'])
df = pd.DataFrame({'a': [1, 2]}, index=idx)
with raises(ValueError):
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
def test_detect_and_parse_datetime_index_invalid_string_error(mlflow):
idx = pd.Index(['bad-format'])
df = pd.DataFrame({'a': [1]}, index=idx)
with raises(ValueError):
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
def test_detect_and_parse_datetime_index_unsupported_type_error(mlflow):
idx = pd.Index([pd.Period('2020-01', freq='M')])
df = pd.DataFrame({'a': [1]}, index=idx)
with raises(ValueError):
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
def test_detect_and_parse_datetime_index_datetime_success(mlflow):
idx = pd.Index([datetime(2020, 1, 1, 0, 0, 0)])
df = pd.DataFrame({'a': [1]}, index=idx)
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
assert out.index[0].endswith('+0000')
def test_detect_and_parse_datetime_index_timestamp_with_tz_success(mlflow):
idx = pd.DatetimeIndex([pd.Timestamp('2020-01-01 00:00:00', tz='UTC')])
df = pd.DataFrame({'a': [1]}, index=idx)
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
assert out.index[0].endswith('+0000')
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
def test_request_transform_success(mock_from_dataframe, mlflow):
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 = MagicMock()
payload.retrieve = MagicMock(return_value=raw)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
response_data = mlflow.request_transform(input_data)
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
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
def test_request_transform_success_without_transform_meta(mock_from_dataframe, mlflow):
ts = pd.Timestamp('2020-01-01', tz='UTC')
raw = pd.DataFrame(
{
'variable': ['v1'],
'timestamp': [ts],
'value': [1.0],
'created_at': [ts],
}
)
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, {})
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
payload = MagicMock()
payload.retrieve = MagicMock(return_value=raw)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
mlflow.request_transform(input_data)
mock_from_dataframe.assert_called_once()
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
def test_request_transform_failure(mock_from_dataframe, mlflow):
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('boom')
data_mock = MagicMock()
payload = MagicMock()
payload.retrieve = MagicMock(return_value=data_mock)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
data_mock.sort_values.return_value = data_mock
data_mock.drop_duplicates.return_value = data_mock
data_mock.pivot.return_value = data_mock
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={'success': False, 'content': ANY},
workflow_metadata=metadata['metadata'],
last_timestamp=payload.last_timestamp,
logger=mlflow.logger,
)
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
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 = MagicMock()
payload.retrieve = MagicMock(return_value=data_mock)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
pred_df.columns = MagicMock()
pred_df.__setitem__ = MagicMock()
response_data = mlflow.request_predict(input_data)
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
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
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
def test_request_predict_success_dataframe_and_meta(mock_to_datetime, mock_from_dataframe, mlflow):
mock_to_datetime.side_effect = lambda x, **kwargs: x
wrapper = MagicMock()
pred_df = pd.DataFrame({'raw': [0.3]})
wrapper.predict.return_value = (pred_df, {'m': 1})
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 = MagicMock()
payload.retrieve = MagicMock(return_value=data_mock)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
mlflow.request_predict(input_data)
assert list(pred_df.columns) == ['prediction', 'response_time']
mlflow.info.assert_any_call("Wrapper predict metadata: {'m': 1}", metadata['metadata'])
mock_from_dataframe.assert_called_once()
@patch(
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
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 = MagicMock()
payload.retrieve = MagicMock(return_value=data_mock)
input_data = {
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
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={'success': False, 'content': ANY},
workflow_metadata=metadata['metadata'],
last_timestamp=payload.last_timestamp,
logger=mlflow.logger,
)
@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_data_success_retrain(
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
):
mock_mkdtemp.return_value = 'tmp'
mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
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],
}
)
pivoted_index = pd.DatetimeIndex([ts])
wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0])
payload = MagicMock()
payload.retrieve = MagicMock(return_value=raw_data)
response = mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {
'target': 'target',
},
}
)
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'
@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_with_payload_data(
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
):
mock_mkdtemp.return_value = 'tmp'
mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
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
ts = pd.Timestamp('2020-01-01', tz='UTC')
raw_data = pd.DataFrame(
{
'variable': ['target', 'f1', 'target', 'f1'],
'timestamp': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
'value': [1.0, 2.0, 3.0, 4.0],
'created_at': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
}
)
pivoted_index = pd.DatetimeIndex([ts, ts + pd.Timedelta(hours=1)])
wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])
payload = MagicMock()
payload.retrieve = MagicMock(return_value=raw_data)
response = mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {'target': 'target'},
}
)
assert response['success'] is True
@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_always_uses_retrain_even_with_full_retrain_flag(
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
):
mock_mkdtemp.return_value = 'tmp'
mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
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
ts = pd.Timestamp('2020-01-01', tz='UTC')
ts_next = ts + pd.Timedelta(days=1)
raw_data = pd.DataFrame(
{
'variable': ['target', 'f1', 'target', 'f1'],
'timestamp': [ts, ts, ts_next, ts_next],
'value': [1.0, 2.0, 3.0, 4.0],
}
)
pivoted_index = pd.DatetimeIndex([ts, ts_next])
wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])
payload = MagicMock()
payload.retrieve = MagicMock(return_value=raw_data)
response = mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {'target': 'target', 'full_retrain': True, 'validation_fraction': 0.5},
}
)
wrapper.retrain.assert_called_once()
wrapper.train.assert_not_called()
assert response['success'] is True
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
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 = MagicMock()
payload.retrieve = MagicMock(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 = mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {'target': 'target'},
}
)
assert response['success'] is False
assert 'retrain failed' in response['message']
def test_retrain_model_data_error(mlflow):
response = mlflow.retrain_model(
{
**metadata,
'model_name': 'test_model',
'model_config': {
'target': 'target',
},
}
)
assert response['success'] is False
assert 'data' in response['message'].lower() or 'loading' in response['message'].lower()
def test_retrain_model_missing_target(mlflow):
ts = pd.Timestamp('2020-01-01', tz='UTC')
raw_data = pd.DataFrame(
{
'variable': ['f1', 'f2'],
'timestamp': [ts, ts],
'value': [1.0, 2.0],
}
)
payload = MagicMock()
payload.retrieve = MagicMock(return_value=raw_data)
response = mlflow.retrain_model(
{
**metadata,
'data': payload,
'model_name': 'test_model',
'model_config': {},
}
)
assert response['success'] is False
assert 'target' in response['message']
def test_retrain_model_data_error_no_minio_repository(mlflow):
mlflow.minio_repository = None
with raises(ValueError) as e:
mlflow.retrain_model(
{
**metadata,
'object_key': 'test_object_key',
'model_name': 'test_model',
'model_config': {
'target': 'target',
},
}
)
assert str(e.value) == 'Minio repository not initialized'
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': {'run_id': 'run-x', 'experiment_id': 'e1'},
'timestamp': 2,
'status': 'success',
}
response = mlflow.update_production_model(input_data)
mlflow.mlflow_repository.promote_to_alias.assert_called_once_with(
model_name='test_model',
version='3',
alias='production',
metadata=metadata['metadata'],
)
assert response['model_name'] == 'test_model'
assert response['version'] == '3'
def test_update_production_model_error(mlflow):
mlflow.mlflow_repository._client.search_model_versions.return_value = []
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 1,
'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
'timestamp': 2,
'status': 'success',
}
try:
mlflow.update_production_model(input_data)
except Exception:
mlflow.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message=ANY,
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('Expected exception')
def _artifact_file_info(path: str) -> MagicMock:
file_info = MagicMock()
file_info.path = path
return file_info
def _retrain_prediction_frame(
index: pd.DatetimeIndex, target: str, values: list[float]
) -> pd.DataFrame:
return pd.DataFrame({target: values}, index=index)
@patch('laborious.activities.mlflow.to_datetime')
def test_get_reference_data_success(mock_to_datetime, mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('evaluation_data.csv'),
]
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},
]
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.object(
mlflow,
'_find_downloaded_csv',
return_value=MagicMock(),
):
result = mlflow.get_reference_data(input_data)
mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
run_id='run1',
artifact_path='evaluation_data.csv',
dst_path='/t',
metadata=metadata['metadata'],
)
mock_reference_data.to_dict.assert_called_once_with(orient='records')
assert result == mock_reference_data.to_dict.return_value
def test_get_reference_data_not_found(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mlflow.mlflow_repository._client.get_model_version_by_alias.side_effect = Exception('missing')
result = mlflow.get_reference_data(input_data)
mlflow.warning.assert_called()
assert result is None
def test_get_reference_data_only_test_data_csv(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('test_data.csv'),
]
mock_reference_data = MagicMock()
mock_reference_data.__getitem__.return_value = MagicMock()
mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
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.object(
mlflow,
'_find_downloaded_csv',
return_value=MagicMock(),
):
result = mlflow.get_reference_data(input_data)
mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
run_id='run1',
artifact_path='test_data.csv',
dst_path='/t',
metadata=metadata['metadata'],
)
assert result == mock_reference_data.to_dict.return_value
def test_get_reference_data_prefers_evaluation_data_when_both_listed(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('test_data.csv'),
_artifact_file_info('evaluation_data.csv'),
]
mock_reference_data = MagicMock()
mock_reference_data.__getitem__.return_value = MagicMock()
mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]
with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
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.object(
mlflow,
'_find_downloaded_csv',
return_value=MagicMock(),
):
result = mlflow.get_reference_data(input_data)
mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
run_id='run1',
artifact_path='evaluation_data.csv',
dst_path='/t',
metadata=metadata['metadata'],
)
assert result == mock_reference_data.to_dict.return_value
def test_get_reference_data_no_candidate_artifacts_returns_none(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('other_artifact.csv'),
]
result = mlflow.get_reference_data(input_data)
mlflow.mlflow_repository.download_artifacts.assert_not_called()
mlflow.warning.assert_called()
assert result is None
def test_get_reference_data_missing_csv_file_returns_none(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('evaluation_data.csv'),
]
with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
with patch('laborious.activities.mlflow.rmtree'):
with patch.object(mlflow, '_find_downloaded_csv', return_value=None):
result = mlflow.get_reference_data(input_data)
mlflow.mlflow_repository.download_artifacts.assert_called_once()
mlflow.warning.assert_called()
assert result is None
def test_get_reference_data_exception(mlflow):
input_data = {
**metadata,
'model_name': 'test_model',
}
mv = MagicMock(run_id='run1')
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
mlflow.mlflow_repository._client.list_artifacts.return_value = [
_artifact_file_info('evaluation_data.csv'),
]
mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
result = mlflow.get_reference_data(input_data)
assert result is None

954
tests/laborious/activities/test_model_metrics.py Normal file
View File

@@ -0,0 +1,954 @@
from unittest.mock import ANY, MagicMock, patch
from pandas import DataFrame, Timestamp
from pytest import fixture, raises
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_model.analytics.drift_analysis import DriftInsufficientDataError
from laborious.activities.model_metrics import ModelMetrics
@fixture
def model_metrics_activity():
model_metrics = ModelMetrics(
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
model_metrics.error = MagicMock()
model_metrics.debug = MagicMock()
model_metrics.info = MagicMock()
model_metrics.warning = MagicMock()
model_metrics.critical = MagicMock()
model_metrics.send_notification = MagicMock()
model_metrics.emit_metric_sync = MagicMock()
model_metrics.get_core_labels = MagicMock(
return_value={
'pod_id': 'test_pod',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
)
model_metrics.observe_lag_sync = MagicMock()
model_metrics.pod_id = 'test_pod'
return model_metrics
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
# Arrange
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': None,
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'invalid',
}
# Act & Assert
try:
model_metrics_activity.calculate_drift(input_data)
except ValueError as e:
assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
model_metrics_activity.error.assert_called_once_with(
'Invalid chunk period: invalid', metadata['metadata']
)
else:
raise AssertionError('Expected ValueError')
def _sample_drift_metrics_df(ts: Timestamp) -> DataFrame:
"""Minimal analyzer-shaped dataframe (univariate row + columns the activity expects)."""
return DataFrame(
{
'timestamp': [ts],
'feature': ['feature1'],
'method': ['ks_test'],
'value': [0.5],
'alert': [False],
'chunk_index': [0],
'chunk_start_date': [ts],
'chunk_end_date': [ts],
'threshold': [0.1],
'drift_type': ['univariate'],
}
)
def test_calculate_drift_with_reference_data(model_metrics_activity):
ts = Timestamp('2023-05-26 11:12:27')
drift_df = _sample_drift_metrics_df(ts)
model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict('list'),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
result = model_metrics_activity.calculate_drift(input_data)
expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
assert result == [
{
'timestamp': expected_timestamp,
'feature': 'feature1',
'method': 'ks_test',
'value': 0.5,
'alert': False,
'chunk_index': 0,
'chunk_start_date': ts.isoformat(),
'chunk_end_date': ts.isoformat(),
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
def test_calculate_drift_without_reference_data(model_metrics_activity):
# Ten rows so int(len * 0.3) >= 1 for the built-in reference slice.
ts_last = Timestamp('2023-05-26 11:12:36')
drift_df = _sample_drift_metrics_df(ts_last)
model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
timestamps = [f'2023-05-26 11:12:{27 + i:02d}' for i in range(10)]
target_data_dict = {
'timestamp': timestamps,
'variable': ['feature1'] * 10,
'value': [float(i) for i in range(10)],
}
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': None,
'target_data': target_data_dict,
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 's',
}
result = model_metrics_activity.calculate_drift(input_data)
expected_timestamp = ts_last.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
assert result == [
{
'timestamp': expected_timestamp,
'feature': 'feature1',
'method': 'ks_test',
'value': 0.5,
'alert': False,
'chunk_index': 0,
'chunk_start_date': ts_last.isoformat(),
'chunk_end_date': ts_last.isoformat(),
'model_id': 'test_model_id',
'accurate': False,
}
]
model_metrics_activity.warning.assert_called()
model_metrics_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
message='Using 30% first rows of target data as reference data',
block='model_metrics',
level=NotificationLevel.WARNING,
attachment_content=ANY,
)
def test_calculate_drift_empty_drift_df(model_metrics_activity):
"""Empty analyzer merge yields no rows and no insufficient-data alert (lib owns that failure mode)."""
ts = Timestamp('2023-05-26 11:12:27')
empty_df = _sample_drift_metrics_df(ts).iloc[0:0]
model_metrics_activity.get_drift_metrics = MagicMock(return_value=empty_df)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict(),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
result = model_metrics_activity.calculate_drift(input_data)
assert result == []
model_metrics_activity.send_notification.assert_not_called()
def test_calculate_drift_empty_after_timestamp_filter(model_metrics_activity):
"""Rows dropped by target-window alignment yield an empty export list, not an insufficient-data error."""
drift_df = _sample_drift_metrics_df(Timestamp('2020-01-01'))
model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict(),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
result = model_metrics_activity.calculate_drift(input_data)
assert result == []
model_metrics_activity.send_notification.assert_not_called()
@patch('laborious.activities.model_metrics.DataFrame')
@patch('laborious.activities.model_metrics.to_datetime')
def test_calculate_drift_drift_insufficient_data_error_from_lib(
mock_to_datetime, mock_dataframe, model_metrics_activity
):
"""``DriftInsufficientDataError`` maps to MODEL_METRICS_DRIFT_INSUFFICIENT_DATA, not GET error."""
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
lib_msg = (
'[MODEL_METRICS_DRIFT_INSUFFICIENT_DATA] Drift analysis produced no time chunks '
"(chunk_period='min', analysis_rows=1)."
)
model_metrics_activity.get_drift_metrics = MagicMock(
side_effect=DriftInsufficientDataError(lib_msg, analysis_rows=1)
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
mock_target_df.reset_index.return_value = mock_target_df
mock_target_df.dropna.return_value = mock_target_df
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict(),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
with raises(DriftInsufficientDataError, match='MODEL_METRICS_DRIFT_INSUFFICIENT_DATA'):
model_metrics_activity.calculate_drift(input_data)
model_metrics_activity.error.assert_called_once_with(lib_msg, metadata['metadata'])
model_metrics_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MODEL_METRICS_DRIFT_INSUFFICIENT_DATA',
message=lib_msg,
block='model_metrics',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
def test_calculate_drift_success_min(model_metrics_activity):
ts = Timestamp('2023-05-26 11:12:27')
drift_df = _sample_drift_metrics_df(ts)
model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict('list'),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
result = model_metrics_activity.calculate_drift(input_data)
expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
assert result == [
{
'timestamp': expected_timestamp,
'feature': 'feature1',
'method': 'ks_test',
'value': 0.5,
'alert': False,
'chunk_index': 0,
'chunk_start_date': ts.isoformat(),
'chunk_end_date': ts.isoformat(),
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
def test_calculate_drift_success_s(model_metrics_activity):
ts = Timestamp('2023-05-26 11:12:27')
drift_df = _sample_drift_metrics_df(ts)
model_metrics_activity.get_drift_metrics = MagicMock(return_value=drift_df)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict('list'),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 's',
}
result = model_metrics_activity.calculate_drift(input_data)
expected_timestamp = ts.tz_localize('UTC').strftime(DATETIME_FORMAT_WITH_TZ)
assert result == [
{
'timestamp': expected_timestamp,
'feature': 'feature1',
'method': 'ks_test',
'value': 0.5,
'alert': False,
'chunk_index': 0,
'chunk_start_date': ts.isoformat(),
'chunk_end_date': ts.isoformat(),
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
@patch('laborious.activities.model_metrics.DataFrame')
@patch('laborious.activities.model_metrics.to_datetime')
def test_calculate_drift_get_drift_metrics_error(
mock_to_datetime, mock_dataframe, model_metrics_activity
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
model_metrics_activity.get_drift_metrics = MagicMock(
side_effect=Exception('Get drift metrics error')
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
mock_target_df.reset_index.return_value = mock_target_df
mock_target_df.dropna.return_value = mock_target_df
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
'model_id': 'test_model_id',
'reference_data': reference_data.to_dict(),
'target_data': {
'timestamp': ['2023-05-26 11:12:27'],
'variable': ['feature1'],
'value': [1.0],
},
'target_name': 'target',
'drift_metrics': ['ks_test'],
'chunk_period': 'min',
}
# Act / Assert
with raises(Exception, match='Get drift metrics error'):
model_metrics_activity.calculate_drift(input_data)
model_metrics_activity.error.assert_called_once_with(
'Error getting drift metrics: Get drift metrics error', metadata['metadata']
)
model_metrics_activity.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
message='Error getting drift metrics: Get drift metrics error',
block='model_metrics',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
@patch('laborious.activities.model_metrics.to_datetime')
@patch('laborious.activities.model_metrics.time.time')
@patch('laborious.activities.model_metrics.DriftAnalysis')
@patch('laborious.activities.model_metrics.metrics')
def test_get_drift_metrics_success(
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
):
# Arrange
mock_time.return_value = 1000.0
mock_drift_df = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'method': ['ks_test'],
'value': [0.5],
'feature': ['feature1'],
}
)
mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
# Act
result = model_metrics_activity.get_drift_metrics(
reference_data=reference_data,
target_data=target_data,
target_name='target',
reference_columns=reference_columns,
drift_metrics=['ks_test'],
chunk_period='min',
metadata=metadata['metadata'],
)
# Assert
assert isinstance(result, DataFrame)
model_metrics_activity.debug.assert_called()
model_metrics_activity.observe_lag_sync.assert_called()
model_metrics_activity.emit_metric_sync.assert_called()
@patch('laborious.activities.model_metrics.to_datetime')
@patch('laborious.activities.model_metrics.time.time')
@patch('laborious.activities.model_metrics.DriftAnalysis')
@patch('laborious.activities.model_metrics.metrics')
def test_get_drift_metrics_univariate_error(
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
):
# Arrange
mock_time.return_value = 1000.0
mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
'Univariate drift error'
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
# Act & Assert
try:
model_metrics_activity.get_drift_metrics(
reference_data=reference_data,
target_data=target_data,
target_name='target',
reference_columns=reference_columns,
drift_metrics=['ks_test'],
chunk_period='min',
metadata=metadata['metadata'],
)
except Exception as e:
assert str(e) == 'Univariate drift error'
model_metrics_activity.error.assert_called_once_with(
'Error detecting univariate drift: Univariate drift error', metadata['metadata']
)
model_metrics_activity.emit_metric_sync.assert_called_with(
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
)
else:
raise AssertionError('Expected Exception')
@patch('laborious.activities.model_metrics.to_datetime')
@patch('laborious.activities.model_metrics.time.time')
@patch('laborious.activities.model_metrics.DriftAnalysis')
@patch('laborious.activities.model_metrics.metrics')
def test_get_drift_metrics_multivariate_error(
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
):
mock_time.return_value = 1000.0
mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.detect_multivariate_drift.side_effect = Exception(
'Multivariate drift error'
)
reference_data = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
)
target_data = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
try:
model_metrics_activity.get_drift_metrics(
reference_data=reference_data,
target_data=target_data,
target_name='target',
reference_columns=reference_columns,
drift_metrics=['ks_test'],
chunk_period='min',
metadata=metadata['metadata'],
)
except Exception as e:
assert str(e) == 'Multivariate drift error'
model_metrics_activity.error.assert_called_once_with(
'Error detecting multivariate drift: Multivariate drift error', metadata['metadata']
)
model_metrics_activity.emit_metric_sync.assert_called_with(
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
)
else:
raise AssertionError('Expected Exception')
@patch('laborious.activities.model_metrics.to_datetime')
@patch('laborious.activities.model_metrics.time.time')
@patch('laborious.activities.model_metrics.DriftAnalysis')
@patch('laborious.activities.model_metrics.metrics')
def test_get_drift_metrics_dataframe_error(
mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity
):
mock_time.return_value = 1000.0
mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.get_drift_metrics_dataframe.side_effect = Exception(
'Dataframe error'
)
reference_data = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
)
target_data = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
try:
model_metrics_activity.get_drift_metrics(
reference_data=reference_data,
target_data=target_data,
target_name='target',
reference_columns=reference_columns,
drift_metrics=['ks_test'],
chunk_period='min',
metadata=metadata['metadata'],
)
except Exception as e:
assert str(e) == 'Dataframe error'
model_metrics_activity.error.assert_called_once_with(
'Error building drift metrics dataframe: Dataframe error', metadata['metadata']
)
model_metrics_activity.emit_metric_sync.assert_called_with(
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
)
else:
raise AssertionError('Expected Exception')
def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse', 'mse', 'mae', 'r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 4
assert 'rmse' in result['metric'].values
assert 'mse' in result['metric'].values
assert 'mae' in result['metric'].values
assert 'r2' in result['metric'].values
assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
metadata['metadata'],
)
model_metrics_activity.debug.assert_called_once()
def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'rmse'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
)
def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['mse'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'mse'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
)
def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['mae'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'mae'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
)
def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'r2'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
# Arrange
# All target values are the same, so ss_tot will be 0
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'r2'
assert result['value'].values[0] == 0.0 # Should return 0.0 when ss_tot == 0
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
# Arrange
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse', 'mae'],
'interval_minutes': 5,
}
# Act
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 2
assert 'rmse' in result['metric'].values
assert 'mae' in result['metric'].values
assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
"Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
)
def test_calculate_simple_metrics_unknown_metric_ignored(model_metrics_activity):
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['unknown_metric', 'rmse'],
'interval_minutes': 5,
}
result = DataFrame(model_metrics_activity.calculate_simple_metrics(input_data))
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'rmse'

720
tests/laborious/activities/test_opc.py Normal file
View File

@@ -0,0 +1,720 @@
from unittest.mock import ANY, MagicMock, call, patch
import pytest
from pandas import DataFrame
from pytest import mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.opc import (
OPC,
OPC_COMMENT_SEPARATOR,
OPC_RECONNECT_IN_PROGRESS_COMMENT,
OPC_SESSION_BAD_COMMENT_PREFIX,
OPC_SESSION_BAD_CONFIDENCE,
OPC_WRITTING_ERROR_CONFIDENCE,
OPC_WRITTING_ERROR_MESSAGE,
_apply_opc_write_error,
)
metadata = {
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test__init__():
servers = {'server1': {'id': 'server1'}}
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
assert opc.opc_servers == servers
assert opc.opc_repository == {}
@patch('laborious.activities.opc.OpcRepository')
@patch('laborious.activities.opc.OPC.send_notification')
def test_init_opc(mock_send_notification, mock_opc_repository):
mock_logger = MagicMock()
mock_metrics_controller = MagicMock()
server1 = MagicMock(
connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
)
server2 = MagicMock(
connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
)
server3 = MagicMock(
connect=MagicMock(
return_value=(
False,
{
'notification_id': 'OPC_CONNECTION_ERROR_server3',
'message': 'Failed to connect to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error',
},
)
),
write_data=MagicMock(return_value=(True, {})),
)
mock_opc_repository.side_effect = [server1, server2, server3]
mock_notification_handler = MagicMock()
servers = {
'server1': {
'server_name': 'server1',
'id': 'server1',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
'server2': {
'server_name': 'server2',
'id': 'server2',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
'server3': {
'server_name': 'server3',
'id': 'server3',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
},
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
opc.init_opc()
assert opc.opc_servers == servers
assert opc.logger == mock_logger
assert opc.notification_handler == mock_notification_handler
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls(
[
call(
opc_id='server1',
server_name='server1',
url='http://localhost:8080',
logger=mock_logger,
server_uri='opc.tcp://localhost:4840',
cert_path='',
private_key_path='',
server_cert_path='',
notification_handler=mock_notification_handler,
reconnection_interval=60,
metrics_controller=mock_metrics_controller,
),
]
)
mock_opc_repository.assert_has_calls(
[
call(
opc_id='server2',
server_name='server2',
url='http://localhost:8080',
logger=mock_logger,
server_uri='opc.tcp://localhost:4840',
cert_path='',
private_key_path='',
server_cert_path='',
notification_handler=mock_notification_handler,
reconnection_interval=60,
metrics_controller=mock_metrics_controller,
)
]
)
server1.connect.assert_called_once()
server2.connect.assert_called_once()
mock_send_notification.assert_has_calls(
[
call(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION',
},
notification_id='OPC_CONNECTION_ERROR_server3',
message='Failed to connect to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
]
)
@pytest.fixture
@patch('laborious.activities.opc.OpcRepository')
def opc(mock_opc_repository):
servers = {
'server1': {
'id': 'server1',
'server_name': 'server1',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
'cert_path': '',
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
}
mock_opc_repository.return_value.write_data = MagicMock(return_value=(True, {}))
mock_opc_repository.return_value.connect = MagicMock(return_value=(True, {}))
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
opc.init_opc()
opc.send_notification = MagicMock()
opc.emit_metric_sync = MagicMock()
return opc
WRITE_DATA_CASES = [
('tag1', 'int', 50),
('tag2', 'float', 50.5),
('tag3', 'bool', True),
('tag4', 'string', 'test'),
]
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
def test_write_data_success(opc, tag, data_type, data):
opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
response_time, error_info = opc.write_data(
server_id='server1',
tag=tag,
data=data,
data_type=data_type,
tag_type='prediction',
metadata=metadata,
)
assert response_time == 0.1
assert error_info is None
opc.opc_repository['server1'].write_data.assert_called_once_with(tag, data, data_type, metadata)
def test_write_data_failed(opc):
opc.opc_repository['server1'].write_data.return_value = (
False,
{
'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
'message': 'Failed to write data to OPC server: Test error',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': 'Test error',
},
)
response_time, error_info = opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
assert response_time is None
assert error_info is not None
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id='OPC_WRITE_DATA_ERROR_server1',
message='Failed to write data to OPC server: Test error',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
try:
opc.write_data(
server_id='server1',
tag='tag1',
data=50,
data_type='int',
tag_type='prediction',
metadata=metadata,
)
except Exception:
opc.send_notification.assert_called_once_with(
metadata=metadata,
notification_id='WRITE_OPC_PREDICTION_ERROR',
message='Error writing data to OPC server: Test error',
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
else:
raise AssertionError('Expected an exception to be raised')
@mark.parametrize(
'error_info,initial_seen,initial_status,initial_reconnect,expected',
[
(None, False, None, False, (False, None, False)),
({}, False, None, False, (False, None, False)),
(
{'opc_error_kind': 'session_bad', 'opc_status': 'BadSessionIdInvalid'},
False,
None,
False,
(True, 'BadSessionIdInvalid', False),
),
(
{'opc_error_kind': 'session_bad', 'opc_status': 'NewStatus'},
True,
'OldStatus',
False,
(True, 'NewStatus', False),
),
(
{'opc_error_kind': 'session_bad'},
True,
'KeptStatus',
False,
(True, 'KeptStatus', False),
),
(
{'opc_error_kind': 'reconnect_in_progress'},
False,
None,
False,
(False, None, True),
),
(
{'opc_error_kind': 'other'},
True,
'Status',
True,
(True, 'Status', True),
),
],
)
def test_apply_opc_write_error(
error_info, initial_seen, initial_status, initial_reconnect, expected
):
result = _apply_opc_write_error(
error_info,
initial_seen,
initial_status,
initial_reconnect,
)
assert result == expected
def test_write_tags_from_config_prediction_success(opc):
opc.write_data = MagicMock(return_value=(0.1, None))
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
tags_config = {'tag1': {'data_type': 'float'}}
response_times, session_bad, opc_status, reconnect = opc._write_tags_from_config(
server_id='server1',
tags_config=tags_config,
data=data,
data_column='prediction',
tag_type='prediction',
log_label='Prediction data',
metadata=metadata['metadata'],
)
assert response_times == {'tag1': 0.1}
assert session_bad is False
assert opc_status is None
assert reconnect is False
opc.write_data.assert_called_once_with(
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction',
metadata=metadata['metadata'],
)
def test_write_tags_from_config_confidence_success(opc):
opc.write_data = MagicMock(return_value=(0.2, None))
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
tags_config = {'tag2': {'data_type': 'float'}}
response_times, session_bad, opc_status, reconnect = opc._write_tags_from_config(
server_id='server1',
tags_config=tags_config,
data=data,
data_column='prediction_confidence',
tag_type='confidence',
log_label='Confidence data',
metadata=metadata['metadata'],
)
assert response_times == {'tag2': 0.2}
assert session_bad is False
assert opc_status is None
assert reconnect is False
opc.write_data.assert_called_once_with(
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence',
metadata=metadata['metadata'],
)
def test_write_tags_from_config_write_failure(opc):
opc.write_data = MagicMock(return_value=(None, {}))
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
response_times, session_bad, opc_status, reconnect = opc._write_tags_from_config(
server_id='server1',
tags_config={'tag1': {'data_type': 'float'}},
data=data,
data_column='prediction',
tag_type='prediction',
log_label='Prediction data',
metadata=metadata['metadata'],
)
assert response_times == {'tag1': None}
assert session_bad is False
assert opc_status is None
assert reconnect is False
def test_write_tags_from_config_session_bad(opc):
opc.write_data = MagicMock(
return_value=(
None,
{
'opc_error_kind': 'session_bad',
'opc_status': 'BadSessionIdInvalid',
},
)
)
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
response_times, session_bad, opc_status, reconnect = opc._write_tags_from_config(
server_id='server1',
tags_config={'tag1': {'data_type': 'float'}},
data=data,
data_column='prediction',
tag_type='prediction',
log_label='Prediction data',
metadata=metadata['metadata'],
)
assert response_times == {'tag1': None}
assert session_bad is True
assert opc_status == 'BadSessionIdInvalid'
assert reconnect is False
def test_write_tags_from_config_reconnect_in_progress(opc):
opc.write_data = MagicMock(
return_value=(
None,
{'opc_error_kind': 'reconnect_in_progress'},
)
)
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
response_times, session_bad, opc_status, reconnect = opc._write_tags_from_config(
server_id='server1',
tags_config={'tag1': {'data_type': 'float'}},
data=data,
data_column='prediction',
tag_type='prediction',
log_label='Prediction data',
metadata=metadata['metadata'],
)
assert response_times == {'tag1': None}
assert session_bad is False
assert opc_status is None
assert reconnect is True
def test_manage_output_tags_success(opc):
opc._write_tags_from_config = MagicMock(
side_effect=[
({'tag1': 0.1}, False, None, False),
({'tag2': 0.1}, False, None, False),
]
)
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
config = {
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
output_data, opc_metrics, session_bad, opc_status, reconnect = opc.manage_output_tags(
server_id='server1',
config=config,
data=data,
metadata=metadata['metadata'],
)
assert output_data is True
assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1}
assert session_bad is False
assert opc_status is None
assert reconnect is False
assert opc._write_tags_from_config.call_count == 2
def test_manage_output_tags_failed(opc):
opc._write_tags_from_config = MagicMock(
side_effect=[
({'tag1': 0.1}, False, None, False),
({'tag2': None}, False, None, False),
]
)
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
config = {
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
output_data, opc_metrics, _, _, _ = opc.manage_output_tags(
server_id='server1',
config=config,
data=data,
metadata=metadata['metadata'],
)
assert output_data is False
assert opc_metrics == {'tag1': 0.1, 'tag2': None}
def test_manage_output_tags_do_nothing(opc):
opc._write_tags_from_config = MagicMock()
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
config = {'_invalid_key': {'tag1': {'data_type': 'float'}}}
output_data, opc_metrics, _, _, _ = opc.manage_output_tags(
server_id='server1',
config=config,
data=data,
metadata=metadata['metadata'],
)
assert output_data is True
assert opc_metrics == {}
opc._write_tags_from_config.assert_not_called()
@patch('laborious.activities.opc.DataFrame')
def test_write_opc_data_success(mock_dataframe, opc):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
},
}
# Act
opc.manage_output_tags = MagicMock(
return_value=(True, {'tag1': 0.1, 'tag2': 0.2}, False, None, False)
)
opc.process_confidence = MagicMock(return_value={'data': 'data'})
output_data, opc_metrics = opc.write_opc_data(input_data)
# Assert
assert output_data == {'data': 'data'}
assert opc_metrics == {'server1': {'tag1': 0.1, 'tag2': 0.2}}
opc.manage_output_tags.assert_called_once_with(
'server1',
input_data['opc_output_config']['server1'],
mock_dataframe.return_value,
metadata['metadata'],
)
opc.process_confidence.assert_called_once_with(
mock_dataframe.return_value,
True,
metadata['metadata'],
session_bad=False,
opc_status=None,
reconnect_in_progress=False,
)
def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
**metadata,
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_servers': ['server1'],
'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
}
# Act
opc.write_opc_data(input_data)
# Assert
opc.opc_repository['server1'].write_data.assert_not_called()
def test_write_opc_data_no_validate_server(opc):
opc.validate_server = MagicMock(return_value=False)
input_data = {
**metadata,
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
},
}
# Act
opc.write_opc_data(input_data)
# Assert
opc.opc_repository['server1'].write_data.assert_not_called()
@mark.parametrize(
'data,success,expected',
[
(DataFrame({'prediction_confidence': [0]}), True, 0),
(DataFrame({'prediction_confidence': [0]}), False, 12),
],
)
def test_process_confidence(opc, data, success, expected):
result = opc.process_confidence(data, success, metadata['metadata'])
assert result['prediction_confidence'][0] == expected
def test_process_confidence_session_bad(opc):
data = DataFrame({'prediction_confidence': [0.9]})
result = opc.process_confidence(
data,
False,
metadata['metadata'],
session_bad=True,
opc_status='BadSessionIdInvalid',
)
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
assert result['comments'][0].startswith(OPC_SESSION_BAD_COMMENT_PREFIX)
assert 'BadSessionIdInvalid' in result['comments'][0]
def test_process_confidence_generic_failure(opc):
data = DataFrame({'prediction_confidence': [0.9]})
result = opc.process_confidence(data, False, metadata['metadata'])
assert result['prediction_confidence'][0] == OPC_WRITTING_ERROR_CONFIDENCE
assert result['comments'][0] == OPC_WRITTING_ERROR_MESSAGE
def test_manage_output_tags_merges_error_flags(opc):
opc._write_tags_from_config = MagicMock(
side_effect=[
({'tag1': None}, True, 'BadSessionIdInvalid', False),
({'tag2': 0.2}, False, None, True),
]
)
data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]})
config = {
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
(
success,
metrics,
session_bad_seen,
opc_status,
reconnect_in_progress,
) = opc.manage_output_tags('server1', config, data, metadata['metadata'])
assert success is False
assert session_bad_seen is True
assert reconnect_in_progress is True
assert opc_status == 'BadSessionIdInvalid'
assert metrics == {'tag1': None, 'tag2': 0.2}
def test_process_confidence_reconnect_in_progress(opc):
data = DataFrame({'prediction_confidence': [0.9]})
result = opc.process_confidence(
data,
False,
metadata['metadata'],
reconnect_in_progress=True,
)
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
assert result['comments'][0] == OPC_RECONNECT_IN_PROGRESS_COMMENT
def test_process_confidence_concatenates_multiple_comments(opc):
data = DataFrame({'prediction_confidence': [0.9]})
session_comment = f'{OPC_SESSION_BAD_COMMENT_PREFIX} BadSessionIdInvalid'
result = opc.process_confidence(
data,
False,
metadata['metadata'],
session_bad=True,
opc_status='BadSessionIdInvalid',
reconnect_in_progress=True,
)
assert result['prediction_confidence'][0] == OPC_SESSION_BAD_CONFIDENCE
assert result['comments'][0] == OPC_COMMENT_SEPARATOR.join(
[session_comment, OPC_RECONNECT_IN_PROGRESS_COMMENT]
)
def test_validate_server(opc):
assert opc.validate_server('server1', metadata) is True
assert opc.validate_server('server2', metadata) is False
def test_close(opc):
repo = opc.opc_repository['server1']
repo.disconnect = MagicMock(return_value=True)
opc.close()
repo.disconnect.assert_called_once()

324
tests/laborious/activities/test_storage.py Normal file
View File

@@ -0,0 +1,324 @@
import datetime
import os
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, raises
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.activities.postgres_sync import Postgres
from laborious.activities.storage import Storage
@fixture(autouse=True)
def _passthrough_from_dict():
with patch(
'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x
):
yield
@fixture(autouse=True)
def _patch_monitoring_shutdown():
"""
Avoid running real async SientiaMonitoring.shutdown when Storage.close runs inside tests.
"""
with patch.object(SientiaMonitoring, 'shutdown') as mock_shutdown:
yield mock_shutdown
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schedule_name': 'test_schedule',
}
}
@fixture
@patch('laborious.activities.storage.MinioRepository')
def storage(mock_minio_repository):
return Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
retention_hours=24,
minio_repository=mock_minio_repository.return_value,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
@patch('laborious.activities.storage.MinioRepository')
def test___init___not_hasattr(mock_minio_repository):
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = MagicMock()
minio_repo = mock_minio_repository.return_value
storage = Storage(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
retention_hours=24,
minio_repository=minio_repo,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
assert isinstance(storage, Postgres)
assert storage.minio_repository is minio_repo
mock_minio_repository.assert_not_called()
@patch('laborious.activities.storage.MinioRepository')
def test___init___none_minio_repository(mock_minio_repository, storage):
storage.minio_repository = None
logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = MagicMock()
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
retention_hours=24,
minio_repository=None,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
assert storage.minio_repository is None
mock_minio_repository.assert_not_called()
@patch('laborious.activities.storage.MinioRepository')
def test___init___done_repository(mock_minio_repository, storage):
storage.__init__(
host='localhost',
port=5432,
user='postgres',
password='postgres',
dbname='postgres',
min_connections=1,
max_connections=10,
retention_hours=24,
minio_repository=mock_minio_repository.return_value,
logger=MagicMock(),
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
mock_minio_repository.assert_not_called()
assert storage.minio_repository is not None
def test_close(storage, _patch_monitoring_shutdown):
storage.minio_repository = MagicMock()
storage.close()
assert storage.minio_repository is None
_patch_monitoring_shutdown.assert_called_once_with(storage)
def test_close_when_minio_repository_already_none(storage, _patch_monitoring_shutdown):
"""Closing without an initialized MinIO repository skips MinIO teardown."""
storage.minio_repository = None
storage.close()
assert storage.minio_repository is None
_patch_monitoring_shutdown.assert_called_once_with(storage)
def test_load_query_with_minio_offload_no_rows(storage):
storage.load_custom_query = MagicMock(return_value=None)
storage_result = {'success': False}
with patch(
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
return_value=storage_result,
) as mock_from_dataframe:
result = storage.load_query_with_minio_offload(
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
)
assert result == storage_result
mock_from_dataframe.assert_called_once()
def test_load_query_with_minio_offload_inline(storage):
storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None}
with patch(
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
return_value=storage_result,
) as mock_from_dataframe:
result = storage.load_query_with_minio_offload(
{
**metadata,
'query': 'SELECT 1',
'model_name': 'my-model',
'key_prefix': 'predictions/s',
}
)
assert result == storage_result
mock_from_dataframe.assert_called_once()
def test_load_query_with_minio_offload_minio(storage):
storage.load_custom_query = MagicMock(return_value=[{'a': 1}])
storage_result = {'success': True, 'data': None, 'object_key': 'object-key'}
with patch(
'laborious.activities.storage.MinioDataFramePayload.from_dataframe',
new_callable=MagicMock,
return_value=storage_result,
) as mock_from_dataframe:
result = storage.load_query_with_minio_offload(
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
)
assert result == storage_result
mock_from_dataframe.assert_called_once()
@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
@patch('laborious.activities.storage.now')
def test_cleanup_minio_objects_expired(mock_now, storage):
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
storage.minio_repository.list_objects = MagicMock(
return_value=[
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
]
)
storage.minio_repository.delete_file = MagicMock()
storage.send_notification = MagicMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
assert result['deleted_count'] == 1
assert result['failed_count'] == 0
deleted_key = (
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
)
assert deleted_key in result['deleted']
assert result['deleted'][deleted_key]['success'] is True
storage.minio_repository.list_objects.assert_called_once_with(
prefix='training_datasets/m',
recursive=True,
metadata=metadata['metadata'],
)
storage.minio_repository.delete_file.assert_called_once_with(
object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
metadata=metadata['metadata'],
)
def test_load_query_with_minio_offload_minio_not_initialized(storage):
storage.minio_repository = None
with raises(ValueError, match='Minio repository not initialized'):
storage.load_query_with_minio_offload({**metadata, 'query': 'SELECT 1', 'model_name': 'm'})
def test_export_payload_to_postgres(storage):
payload = MagicMock()
payload.retrieve = MagicMock(return_value=MagicMock())
storage.export_data_to_postgres = MagicMock(return_value={'success': True})
result = storage.export_payload_to_postgres(
{**metadata, 'data': payload, 'schema': 'public', 'table': 't'}
)
payload.retrieve.assert_called_once_with(storage.minio_repository, metadata['metadata'])
storage.export_data_to_postgres.assert_called_once()
assert result == {'success': True}
def test_cleanup_minio_objects_expired_minio_not_initialized(storage):
storage.minio_repository = None
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'test'
with raises(ValueError, match='Minio repository not initialized'):
storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
@patch('laborious.activities.storage.now')
def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage):
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
storage.minio_repository.list_objects = MagicMock(
return_value=['some/random/key-without-timestamp.parquet']
)
storage.minio_repository.delete_file = MagicMock()
storage.send_notification = MagicMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'test'
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
assert result['deleted_count'] == 0
assert result['failed_count'] == 0
storage.minio_repository.delete_file.assert_not_called()
@patch('laborious.activities.storage.now')
def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage):
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet'
storage.minio_repository.list_objects = MagicMock(return_value=[old_key])
storage.minio_repository.delete_file = MagicMock(side_effect=Exception('delete error'))
storage.send_notification = MagicMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
assert result['deleted_count'] == 0
assert result['failed_count'] == 1
assert old_key in result['failed']
assert result['failed'][old_key]['success'] is False
assert result['failed'][old_key]['message'] == 'delete error'
@patch('laborious.activities.storage.now')
def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage):
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
storage.minio_repository.list_objects = MagicMock(side_effect=Exception('list error'))
storage.send_notification = MagicMock()
storage.error = MagicMock()
data_mock = MagicMock()
data_mock.cleanup_prefix.return_value = 'training_datasets/m'
result = storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
assert result['deleted_count'] == 0
assert result['failed_count'] == 0
storage.send_notification.assert_called_once_with(
metadata=metadata['metadata'],
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
message='Error cleaning up MinIO objects: list error',
block='cleanup_minio_objects_expired',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)
storage.error.assert_called_once()