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

SIENTIAPDE-1646 Sync full repo from local release/SIENTIAPDE-1646 (272e02d)

Browse Source
This commit is contained in:
vitor-aignosi
2026-07-15 17:54:58 -03:00
parent c76f24c13a
commit be39833dbd
71 changed files with 15501 additions and 2919 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,18 +1,32 @@
from pytest import mark
from unittest.mock import patch, MagicMock, ANY
from sientia_do.temporal.activities.postgres import Postgres
from unittest.mock import ANY, MagicMock, patch
from laborious.activities.activities import Activities
from laborious.activities.mlflow import MLFlow
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.Postgres.__init__')
@patch('laborious.activities.activities.Storage.__init__')
@patch('laborious.activities.activities.MLFlow.__init__')
@patch('laborious.activities.activities.OPC.__init__')
@patch('laborious.activities.activities.Gates.__init__')
def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_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,
@@ -20,20 +34,31 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
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'
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
@@ -41,19 +66,24 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
plugin_store=plugin_store,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
mlflow_repository=mlflow_repository,
)
assert isinstance(activities, Activities)
assert isinstance(activities, Postgres)
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_postgres_init.assert_called_once_with(
mock_storage_init.assert_called_once_with(
ANY,
host=postgres_config['host'],
port=postgres_config['port'],
@@ -62,40 +92,85 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
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
notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
)
mock_mlflow_init.assert_called_once_with(
ANY,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
mlflow_repository=mlflow_repository,
plugin_store=plugin_store,
minio_repository=mock_minio_repository.return_value,
logger=logger,
notification_handler=notification_handler
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
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
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'],
)
@mark.asyncio
@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
@patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
@patch('laborious.activities.activities.OPC', return_value=MagicMock())
async def test_shutdown(mock_opc_init,
_mock_mlflow_init, mock_postgres_init):
@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,
@@ -103,20 +178,31 @@ async def test_shutdown(mock_opc_init,
'password': 'postgres',
'dbname': 'postgres',
'min_connections': 1,
'max_connections': 10
'max_connections': 10,
}
mlflow_config = {
'host': 'localhost',
'port': 5000,
'username': 'mlflow',
'password': 'mlflow'
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'
'group_id': 'test-group',
}
pi_web_api_config = {
'base_url': 'https://test-pi-server.com',
'auth_type': 'bearer',
'auth_token': 'test_token',
}
logger = MagicMock()
@@ -124,12 +210,90 @@ async def test_shutdown(mock_opc_init,
activities = Activities(
postgres_config=postgres_config,
mlflow_config=mlflow_config,
plugin_store=plugin_store,
minio_config=minio_config,
opc_config=opc_config,
pi_web_api_config=pi_web_api_config,
logger=logger,
notification_handler=notification_handler
notification_handler=notification_handler,
mlflow_repository=mlflow_repository,
)
await activities.shutdown()
mock_opc_init.shutdown.assert_called_once()
mock_postgres_init.close.assert_called_once()
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']
)

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

File diff suppressed because it is too large Load Diff

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

File diff suppressed because it is too large Load Diff

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'

743
tests/laborious/activities/test_opc.py
View File

@@ -1,62 +1,75 @@
from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
from unittest.mock import ANY, MagicMock, call, patch
import pytest
from pandas import DataFrame
from pytest import fixture, mark
import pytest_asyncio
from pytest import mark
from sientia_do.notifications.models import NotificationLevel
from laborious.activities.opc import OPC
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",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test__init__():
servers = {
'server1': 'config'
}
servers = {'server1': {'id': 'server1'}}
opc = OPC(
opc_servers=servers,
logger=MagicMock(),
notification_handler=MagicMock()
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
assert opc.opc_servers == servers
assert opc.opc_repository == {}
@mark.asyncio
@patch("laborious.activities.opc.OpcRepository")
@patch("laborious.activities.opc.OPC.send_notification")
async def test_init_opc(mock_send_notification, mock_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=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
)
server2 = MagicMock(
connect=AsyncMock(return_value=(True, {})),
write_data=AsyncMock(return_value=(True, {}))
connect=MagicMock(return_value=(True, {})), write_data=MagicMock(return_value=(True, {}))
)
server3 = MagicMock(
connect=AsyncMock(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=AsyncMock(return_value=(True, {}))
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',
@@ -66,6 +79,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
'reconnection_interval': 60,
},
'server2': {
'server_name': 'server2',
'id': 'server2',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
@@ -75,6 +89,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
'reconnection_interval': 60,
},
'server3': {
'server_name': 'server3',
'id': 'server3',
'url': 'http://localhost:8080',
'server_uri': 'opc.tcp://localhost:4840',
@@ -82,14 +97,15 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
'private_key_path': '',
'server_cert_path': '',
'reconnection_interval': 60,
}
},
}
opc = OPC(
opc_servers=servers,
logger=mock_logger,
notification_handler=mock_notification_handler
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
await opc.init_opc()
opc.init_opc()
assert opc.opc_servers == servers
assert opc.logger == mock_logger
@@ -97,61 +113,70 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
assert opc.opc_repository['server1'] == server1
assert opc.opc_repository['server2'] == server2
mock_opc_repository.assert_has_calls([
call(
id="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,
pod_id='localhost'
),
])
mock_opc_repository.assert_has_calls([
call(
id="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,
pod_id='localhost'
)
])
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
)
])
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_asyncio.fixture
@patch("laborious.activities.opc.OpcRepository")
async def opc(mock_opc_repository):
@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': '',
@@ -161,19 +186,17 @@ async def opc(mock_opc_repository):
}
}
mock_opc_repository.return_value.write_data = AsyncMock(
return_value=(True, {})
)
mock_opc_repository.return_value.connect = AsyncMock(
return_value=(True, {})
)
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()
notification_handler=MagicMock(),
metrics_controller=MagicMock(),
)
await opc.init_opc()
opc.init_opc()
opc.send_notification = MagicMock()
opc.emit_metric_sync = MagicMock()
return opc
@@ -186,184 +209,512 @@ WRITE_DATA_CASES = [
@mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
@mark.asyncio
async def test_write_data_success(opc, tag, data_type, data):
result = await opc.write_data(server_id='server1', tag=tag, data=data,
data_type=data_type, tag_type='prediction', metadata=metadata)
assert result is True
opc.opc_repository['server1'].write_data.assert_called_once_with(
tag, data, data_type, opc.logger, metadata)
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)
@mark.asyncio
async 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'
})
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',
},
)
result = await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
assert result is False
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",
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
attachment_content=ANY,
)
@mark.asyncio
async def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception(
"Test error")
def test_write_data_exception(opc):
opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
try:
await opc.write_data(server_id='server1', tag='tag1', data=50,
data_type='int', tag_type='prediction', metadata=metadata)
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",
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
attachment_content=ANY,
)
else:
assert False, "Expected an exception to be raised"
raise AssertionError('Expected an exception to be raised')
@mark.asyncio
async def test_write_opc_data_success(opc):
@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]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
}
},
}
# Act
opc.write_data = AsyncMock(return_value=True)
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 = await opc.write_opc_data(input_data)
output_data, opc_metrics = opc.write_opc_data(input_data)
# Assert
assert output == {'data': 'data'}
opc.write_data.assert_has_calls([
call(
server_id='server1',
tag='tag1',
data=0.75,
data_type='float',
tag_type='prediction',
metadata=metadata['metadata']
)])
opc.write_data.assert_has_calls([
call(
server_id='server1',
tag='tag2',
data=0.95,
data_type='float',
tag_type='confidence',
metadata=metadata['metadata']
)
])
assert opc.write_data.call_count == 2
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,
)
@mark.asyncio
async def test_write_opc_data_empty_config(opc):
def test_write_opc_data_empty_config(opc):
# Arrange
input_data = {
**metadata,
'data': {
'prediction': [0.75],
'prediction_confidence': [0.95]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_servers': ['server1'],
'opc_output_config': {
'server1': {
'prediction_tags': {},
'confidence_tags': {}
}
}
'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
}
# Act
await opc.write_opc_data(input_data)
opc.write_opc_data(input_data)
# Assert
opc.opc_repository['server1'].write_data.assert_not_called()
@mark.asyncio
async def test_write_opc_data_no_validate_server(opc):
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]
},
'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
'opc_output_config': {
'server1': {
'prediction_tags': {
'tag1': {'data_type': 'float'}
},
'confidence_tags': {
'tag2': {'data_type': 'float'}
}
'prediction_tags': {'tag1': {'data_type': 'float'}},
'confidence_tags': {'tag2': {'data_type': 'float'}},
}
}
},
}
# Act
await opc.write_opc_data(input_data)
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),
])
@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):
# Act
result = opc.process_confidence(data, success, metadata)
# Assert
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
@mark.asyncio
async def test_shutdown(opc):
opc.opc_repository['server1'].disconnect = AsyncMock(return_value=True)
await opc.shutdown()
opc.opc_repository['server1'].disconnect.assert_called_once()
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()

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

@@ -1,23 +1,36 @@
from pandas import DataFrame
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
filter_empty_data
)
def test_filter_specific_variables_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']}) is False
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
config={'variables': ['variable2']},
)
is False
)
def test_filter_specific_variables_null_values_with_empty_data():
assert (
filter_specific_variables_null_values(DataFrame(), config={'variables': ['variable2']})
is False
)
def test_filter_specific_variables_null_values_with_null_values():
assert filter_specific_variables_null_values(
DataFrame(
{'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']}) is True
assert (
filter_specific_variables_null_values(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
config={'variables': ['variable2']},
)
is True
)
def test_filter_empty_data():
@@ -25,6 +38,7 @@ def test_filter_empty_data():
def test_filter_empty_data_with_data():
assert filter_empty_data(
DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
{}) is False
assert (
filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
is False
)

11
tests/laborious/utils/filters/test_mlflow_filters.py
View File

@@ -1,22 +1,23 @@
from pandas import DataFrame
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
def test_api_error_filter_invalid_response():
assert api_error_filter(None, {}) == True # NOSONAR
assert api_error_filter(None, {}) is True # NOSONAR
def test_api_error_filter_valid_response_fail():
assert api_error_filter({'success': False}, {}) == True
assert api_error_filter({'success': False}, {}) is True
def test_api_error_filter_valid_response_success():
assert api_error_filter({'success': True}, {}) == False
assert api_error_filter({'success': True}, {}) is False
def test_nan_values_filter_all_nan_values():
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) is True
def test_nan_values_filter_no_nan_values():
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) is False

278
tests/laborious/utils/models/test_minio_dataframe_payload.py Normal file
View File

@@ -0,0 +1,278 @@
from datetime import datetime
from io import BytesIO
from unittest.mock import MagicMock, patch
from pandas import DataFrame
from laborious.utils.models.minio_dataframe_payload import (
MinioDataFramePayload,
_build_object_key,
)
def test_parse_object_timestamp_hyphenated_model():
key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
ts = MinioDataFramePayload.parse_object_timestamp(key)
assert ts == datetime(2024, 6, 15, 10, 30, 45)
def test_parse_object_timestamp_transform():
key = 'p/m-transform-2024-01-02_03-04-05.parquet'
ts = MinioDataFramePayload.parse_object_timestamp(key)
assert ts == datetime(2024, 1, 2, 3, 4, 5)
def test_parse_object_timestamp_invalid():
assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
def test_estimate_size_bytes_returns_positive_for_nonempty_frame():
df = DataFrame({'a': [1, 2]})
size = MinioDataFramePayload.estimate_size_bytes(df)
assert isinstance(size, int)
assert size > 0
def test_cleanup_prefix_when_offloaded_returns_object_prefix():
payload = MinioDataFramePayload(
last_timestamp='t',
data=None,
object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
object_prefix='training_datasets/m',
)
assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m'
def test_cleanup_prefix_when_inline_returns_none():
payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
assert MinioDataFramePayload.cleanup_prefix(payload) is None
def test_has_data_true_when_object_key_set():
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
assert payload.has_data() is True
def test_retrieve_inline_dict_as_dataframe():
payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
minio = MagicMock()
out = payload.retrieve(minio, {'metadata': {}})
assert list(out.columns) == ['a']
minio.download_file.assert_not_called()
def test_retrieve_downloads_parquet_when_offloaded():
source = DataFrame({'a': [1, 2]})
buf = BytesIO()
source.to_parquet(buf, engine='pyarrow', index=True)
file_bytes = buf.getvalue()
payload = MinioDataFramePayload(
last_timestamp='t',
data=None,
object_key='training_datasets/m/f.parquet',
object_prefix='training_datasets/m',
)
minio = MagicMock()
minio.download_file = MagicMock(return_value=file_bytes)
out = payload.retrieve(minio, {'metadata': {}})
minio.download_file.assert_called_once_with(
object_name='training_datasets/m/f.parquet',
metadata={'metadata': {}},
)
assert list(out.columns) == ['a']
def test_build_object_key():
key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
assert prefix == 'prediction_datasets/my-model'
def test_build_object_key_strips_slashes():
key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
assert prefix == 'prediction_datasets/my-model'
assert key.startswith('prediction_datasets/my-model/')
def test_estimate_size_bytes_fallback():
df = DataFrame({'a': [1, 2]})
with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')):
size = MinioDataFramePayload.estimate_size_bytes(df)
assert isinstance(size, int)
assert size > 0
def test_parse_object_timestamp_bad_datetime():
key = 'p/m-initial-9999-99-99_99-99-99.parquet'
assert MinioDataFramePayload.parse_object_timestamp(key) is None
def test_retrieve_empty_when_no_data():
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
minio = MagicMock()
out = payload.retrieve(minio, {})
assert out.empty
minio.download_file.assert_not_called()
@patch('laborious.utils.models.minio_dataframe_payload.now')
def test_from_dataframe_none(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = MagicMock()
result = MinioDataFramePayload.from_dataframe(
dataframe=None,
minio_repo=minio,
model_name='m',
operation='initial',
status={'success': False, 'message': 'no data'},
)
assert result.data is None
assert result.status == {'success': False, 'message': 'no data'}
assert result.object_key is None
@patch('laborious.utils.models.minio_dataframe_payload.now')
def test_from_dataframe_empty(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = MagicMock()
mock_df = MagicMock()
mock_df.__bool__ = MagicMock(return_value=True)
mock_df.empty = True
result = MinioDataFramePayload.from_dataframe(
dataframe=mock_df,
minio_repo=minio,
model_name='m',
operation='initial',
)
assert result.data is None
assert result.object_key is None
def _mock_dataframe(data_dict, timestamp_values=None):
"""Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
mock_df = MagicMock()
mock_df.__bool__ = MagicMock(return_value=True)
mock_df.empty = False
if timestamp_values is None:
timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
ts_col = MagicMock()
ts_col.values.tolist.return_value = timestamp_values
mock_df.__getitem__ = MagicMock(return_value=ts_col)
mock_df.to_dict.return_value = data_dict
buf = BytesIO()
DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
return mock_df
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
def test_from_dataframe_inline():
minio = MagicMock()
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
result = MinioDataFramePayload.from_dataframe(
dataframe=df,
minio_repo=minio,
model_name='m',
operation='initial',
)
assert result.data is not None
assert result.object_key is None
assert result.last_timestamp == '2024-01-01'
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
def test_from_dataframe_inline_uses_provided_last_timestamp():
minio = MagicMock()
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
result = MinioDataFramePayload.from_dataframe(
dataframe=df,
minio_repo=minio,
model_name='m',
operation='initial',
last_timestamp='2024-01-02',
)
assert result.last_timestamp == '2024-01-02'
@patch('laborious.utils.models.minio_dataframe_payload.now')
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
def test_from_dataframe_offloaded(mock_now):
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
minio = MagicMock()
minio.upload_file = MagicMock(return_value={'minio_object_name': 'full/key.parquet'})
minio.bucket = 'test-bucket'
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
result = MinioDataFramePayload.from_dataframe(
dataframe=df,
minio_repo=minio,
model_name='m',
operation='initial',
workflow_metadata={'wf': 'data'},
)
assert result.data is None
assert result.object_key == 'full/key.parquet'
assert result.bucket == 'test-bucket'
assert result.uri == 's3://test-bucket/full/key.parquet'
minio.upload_file.assert_called_once()
def test_from_dict_inline():
raw = {
'last_timestamp': '2024-01-01T00:00:00+00:00',
'status': None,
'data': {'col1': {0: 'val1'}},
'bucket': None,
'object_key': None,
'object_prefix': None,
'uri': None,
}
payload = MinioDataFramePayload.from_dict(raw)
assert isinstance(payload, MinioDataFramePayload)
assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
assert payload.data == {'col1': {0: 'val1'}}
assert payload.object_key is None
def test_from_dict_offloaded():
raw = {
'last_timestamp': '2024-06-15T10:30:45+00:00',
'status': {'success': True},
'data': None,
'bucket': 'my-bucket',
'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
'object_prefix': 'training_datasets/model',
'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
}
payload = MinioDataFramePayload.from_dict(raw)
assert isinstance(payload, MinioDataFramePayload)
assert payload.data is None
assert payload.bucket == 'my-bucket'
assert payload.object_key == raw['object_key']
assert payload.object_prefix == 'training_datasets/model'
assert payload.uri == raw['uri']
assert payload.status == {'success': True}
def test_from_dict_minimal_keys():
raw = {'last_timestamp': '2024-01-01'}
payload = MinioDataFramePayload.from_dict(raw)
assert payload.last_timestamp == '2024-01-01'
assert payload.data is None
assert payload.bucket is None
assert payload.object_key is None
def test_from_dict_passthrough_existing_instance():
original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
result = MinioDataFramePayload.from_dict(original)
assert result is original
def test_debug_with_logger_calls_custom_debug():
logger = MagicMock()
MinioDataFramePayload._debug(logger, 'msg', {'a': 1})
logger.custom_debug.assert_called_once_with('msg', {'a': 1})

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

@@ -1,502 +0,0 @@
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
from pandas import DataFrame
import pytest
from datetime import datetime, timezone
from pandas import Timestamp
from laborious.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch('laborious.utils.repository.model_repository.ModelServing',
autospec=True) as mock_model_serving:
mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000',
username='admin',
password='admin',
logger=MagicMock()
)
return repo
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
},
}
class Any:
pass
invalid_cases = [
(
{
'value': {
'2024-01-01 12:00:00': 1,
2024: 2
}
}
),
(
{
'value': {
'2024-01-01': 1,
'2024-01-02': 2
}
}
),
(
{
'value': {
Any(): 1,
Any(): 2
}
}
)
]
@pytest.mark.parametrize("data", invalid_cases)
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
input_data = DataFrame(
data
)
with pytest.raises(ValueError) as e:
mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S"
valid_cases = [
(
{
'value': {
'2024-01-01 12:00:00+0000': 1,
'2024-01-02 12:00:00+0000': 2
}
}, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000']
),
(
{
'value': {
datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
}
}, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000']
),
(
{
'value': {
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
}
}, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000']
),
]
@pytest.mark.parametrize("data,expected", valid_cases)
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
input_data = DataFrame(data)
response = mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
assert response.index.tolist() == expected
def test_transform_success(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata'])
assert output == {
'success': True,
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value
}
def test_transform_error(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
'error')
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
def test_predict_success(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
)
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
assert output['success'] is True
assert output['content'] == {
'prediction': {
'index_1': 2,
'index_2': 3
}, 'response_time': {
'index_1': ANY,
'index_2': ANY
}
}
def test_predict_error(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(
side_effect=Exception('error')
)
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
mlflow.get_run.return_value = MagicMock(
info=MagicMock(
experiment_id='0',
)
)
mlflow.get_experiment.return_value = MagicMock()
mlflow.get_experiment.return_value.name = 'test'
output = mlflow_repository.get_experiment_by_run_id('0')
assert output == 'test'
mlflow.get_run.assert_called_once_with('0')
mlflow.get_experiment.assert_called_once_with('0')
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_next_run_name(mlflow, mlflow_repository):
mlflow.search_runs.return_value = [1, 2, 3]
output = mlflow_repository.get_next_run_name('run')
assert output == 'run-4'
mlflow.search_runs.assert_called_once_with(
experiment_names=['run'],
order_by=['start_time desc'],
)
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
output = mlflow_repository.get_experiment('test')
assert output == 0
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_error(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = None
try:
mlflow_repository.get_experiment('test')
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run(mlflow, mlflow_repository):
mlflow.search_runs.return_value = DataFrame({
'params.retrain': ['True', 'False', 'True', 'False'],
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
'run_id': ['0', '1', '2', '3'],
})
output = mlflow_repository.get_experiment_last_run(0)
mlflow.search_runs.assert_called_once_with(
experiment_ids=[0],
filter_string="",
output_format="pandas",
)
assert output == '2'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_get_experiment_last_run_error(mlflow, mlflow_repository):
mlflow.search_runs.return_value = []
try:
mlflow_repository.get_experiment_last_run(0)
except ValueError as e:
assert str(e) == 'Runs is not a pandas DataFrame'
else:
assert False
@patch('laborious.utils.repository.model_repository.mlflow.sklearn')
@patch('laborious.utils.repository.model_repository.mlflow.set_experiment')
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
mlflow_repository.model_serving.get_model_run_id = MagicMock(
return_value='0')
mlflow_repository.model_serving.get_model_uri = MagicMock(
return_value='test')
mlflow_repository.get_experiment_by_run_id = MagicMock()
data_model_mock = MagicMock()
prediction_model_mock = MagicMock()
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
data_model_mock.fit.return_value = data_model_mock
data_model_mock.predict.return_value = DataFrame({
'x': [10, 20, 30],
})
data_model_mock.target_variable = 'y'
prediction_model_mock.fit.return_value = prediction_model_mock
data = DataFrame({
'x': [1, 2, 3],
'y': [4, 5, 6]
})
output = mlflow_repository.create_model_experiment('test', data)
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
'test', stage='Production')
mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
'0', prediction=False)
sklearn.load_model.assert_has_calls([
call(mlflow_repository.model_serving.get_model_uri.return_value),
call("models:/test/production"),
])
assert sklearn.load_model.call_count == 2
data_model_mock.fit.assert_called_once_with(data)
data_model_mock.predict.assert_called_once_with(data)
fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals(
DataFrame({
'x': [10, 20, 30],
'y': [4, 5, 6],
})
)
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
set_experiment.assert_called_once_with(
mlflow_repository.get_experiment_by_run_id.return_value
)
assert output == (prediction_model_mock,
data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value)
@patch('laborious.utils.repository.model_repository.mlflow.start_run')
@patch('laborious.utils.repository.model_repository.mlflow.log_param')
@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model')
@patch('laborious.utils.repository.model_repository.mlflow.log_artifact')
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
prediction_model_mock = MagicMock()
data_model_mock = MagicMock()
experiment = 'test'
model_name = 'test'
data = MagicMock()
mlflow_repository.get_next_run_name = MagicMock(
return_value='test-1')
run = MagicMock()
start_run.__enter__.return_value = run
output = mlflow_repository.perform_model_retrain(
prediction_model_mock, data_model_mock, experiment, model_name, data)
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
start_run.assert_called_once_with(
run_name='test-1', description='Retrain model test with new data')
log_model.assert_has_calls([
call(data_model_mock, "data_model"),
call(prediction_model_mock, "prediction_model"),
])
data.to_csv.assert_called_once_with(
"temp/raw_data_test.csv", index=True)
log_artifact.assert_called_once_with(
"temp/raw_data_test.csv")
log_param.assert_has_calls([
call("retrain", True),
])
assert output == ("Model retrained successfully", experiment)
def test_retrain_model(mlflow_repository):
data = MagicMock()
model_name = 'test'
mlflow_repository.create_model_experiment = MagicMock(
return_value=('data_model', 'prediction_model', '0'))
mlflow_repository.perform_model_retrain = MagicMock(
return_value='Model retrained successfully')
output = mlflow_repository.retrain_model(data, model_name)
mlflow_repository.create_model_experiment.assert_called_once_with(
model_name, data)
mlflow_repository.perform_model_retrain.assert_called_once_with(
'data_model', 'prediction_model', '0', model_name, data)
assert output == 'Model retrained successfully'
@patch('laborious.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
client_mock = MagicMock()
mlflow.tracking.MlflowClient.return_value = client_mock
client_mock.get_registered_model.return_value = MagicMock(
latest_versions=[
MagicMock(version='1'),
MagicMock(version='2'),
MagicMock(version='3'),
]
)
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'test',
)
mlflow.tracking.MlflowClient.assert_called_once()
client_mock.get_registered_model.assert_called_once_with('test')
client_mock.transition_model_version_stage.assert_called_once_with(
name='test',
version='3',
stage='Production',
archive_existing_versions=True,
)
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
}
@patch('laborious.utils.repository.model_repository.mlflow')
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
mlflow.tracking.MlflowClient.return_value = MagicMock(
get_registered_model=MagicMock(
return_value=MagicMock(
latest_versions={}
)
)
)
try:
mlflow_repository.update_production_model_by_run_id('0', 'test')
except Exception as e:
assert str(e) == 'Model versions is not a list'
else:
assert False
def test_update_production_model(mlflow_repository):
connector = mlflow_repository
with patch.object(connector, 'get_experiment',
return_value='0') as get_experiment:
with patch.object(connector, 'get_experiment_last_run',
return_value='2') as get_experiment_last_run:
with patch.object(connector, 'update_production_model_by_run_id',
return_value={'model_name': 'test', 'version': '3',
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
output = connector.update_production_model('0', 'test')
get_experiment.assert_called_once_with('0')
get_experiment_last_run.assert_called_once_with('0')
update_production_model_by_run_id.assert_called_once_with(
'2', 'test')
assert output == {
'model_name': 'test',
'version': '3',
'mlflow_run_id': '0',
'mlflow_experiment_id': '0',
}

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

@@ -1,9 +1,20 @@
import pytest
from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from laborious.utils.repository.opc_repository import OpcRepository
from sientia_do.notifications.models import NotificationLevel
import concurrent.futures
import json
from datetime import datetime
from unittest.mock import ANY, MagicMock, Mock, patch
import pytest
from asyncua.crypto import security_policies
from asyncua.ua.uaerrors import BadNodeIdUnknown, BadSessionIdInvalid
from sientia_do.notifications.models import NotificationLevel
from laborious.utils.repository.opc_repository import (
OpcClientAlreadyExistsError,
OpcClientNotInitializedError,
OpcRepository,
OpcSessionAlreadyConnectedError,
is_reconnectable_opcua_bad,
)
@pytest.fixture
@@ -13,376 +24,561 @@ def mock_logger():
@pytest.fixture
def opc_repository(mock_logger):
return OpcRepository(
id="test_repo",
url="opc.tcp://localhost:4840",
repository = OpcRepository(
opc_id='test_repo',
server_name='test_server',
url='opc.tcp://localhost:4840',
logger=mock_logger,
notification_handler=Mock(),
reconnection_interval=60,
server_uri="urn:test:server",
cert_path="/path/to/cert.pem",
private_key_path="/path/to/key.pem",
server_cert_path="/path/to/server_cert.pem"
server_uri='urn:test:server',
cert_path='/path/to/cert.pem',
private_key_path='/path/to/key.pem',
server_cert_path='/path/to/server_cert.pem',
metrics_controller=MagicMock(),
)
repository.disconnection_interval = 0.1
repository.send_notification = MagicMock()
repository.send_notification = MagicMock()
repository.emit_metric_sync = MagicMock()
repository.info = MagicMock()
repository.error = MagicMock()
repository.warning = MagicMock()
repository.debug = MagicMock()
repository._session_ready.set()
return repository
@pytest.fixture
def mock_client():
with patch('laborious.utils.repository.opc_repository.Client') as mock:
client_instance = AsyncMock()
client_instance = MagicMock()
aio = MagicMock()
client_instance.aio_obj = aio
aio.uaclient = MagicMock()
aio.uaclient.protocol = MagicMock(state='closed')
aio.session_timeout = 600_000
aio.secure_channel_timeout = 600_000
mock.return_value = client_instance
yield client_instance
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
def test_init(opc_repository):
assert opc_repository.id == "test_repo"
assert opc_repository.url == "opc.tcp://localhost:4840"
assert opc_repository.server_uri == "urn:test:server"
assert opc_repository.cert_path == "/path/to/cert.pem"
assert opc_repository.private_key_path == "/path/to/key.pem"
assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
assert opc_repository.id == 'test_repo'
assert opc_repository.server_name == 'test_server'
assert opc_repository.url == 'opc.tcp://localhost:4840'
assert opc_repository.server_uri == 'urn:test:server'
assert opc_repository.cert_path == '/path/to/cert.pem'
assert opc_repository.private_key_path == '/path/to/key.pem'
assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
assert opc_repository.reconnection_interval == 60
assert opc_repository.client is None
assert opc_repository.last_reconnection_time is None
assert opc_repository.error_count == 0
@pytest.mark.asyncio
async def test_set_security(opc_repository, mock_client):
def test_set_security(opc_repository, mock_client):
opc_repository.client = mock_client
await opc_repository.set_security()
opc_repository.set_security()
mock_client.application_uri = "urn:test:server"
mock_client.application_uri = 'urn:test:server'
mock_client.set_security.assert_called_once_with(
SecurityPolicyBasic256,
certificate="/path/to/cert.pem",
private_key="/path/to/key.pem",
server_certificate="/path/to/server_cert.pem"
security_policies.SecurityPolicyBasic256,
'/path/to/cert.pem',
'/path/to/key.pem',
None,
'/path/to/server_cert.pem',
)
assert mock_client.secure_channel_timeout == 10000000
assert mock_client.session_timeout == 10000000
assert mock_client.aio_obj.secure_channel_timeout == 600_000
assert mock_client.aio_obj.session_timeout == 600_000
@pytest.mark.asyncio
async def test_set_security_missing_certificates(opc_repository):
def test_set_security_missing_certificates(opc_repository):
opc_repository.cert_path = None
opc_repository.private_key_path = None
try:
await opc_repository.set_security()
opc_repository.set_security()
except ValueError as e:
assert str(
e) == "Certificate and private key paths must be provided for secure connection."
assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
@pytest.mark.asyncio
async def test_connect_with_security(opc_repository, mock_client):
opc_repository.try_connect = AsyncMock(return_value=(True, {}))
result = await opc_repository.connect()
def test_set_security_missing_client(opc_repository):
opc_repository.client = None
try:
opc_repository.set_security()
except ValueError as e:
assert str(e) == 'Client must be initialized before setting security'
opc_repository.try_connect.assert_called_once()
assert opc_repository.client == mock_client
def test_connect_with_security(opc_repository, mock_client):
opc_repository._create_client = MagicMock()
opc_repository._open_session = MagicMock(return_value=(True, {}))
result = opc_repository.connect()
opc_repository._create_client.assert_called_once()
opc_repository._open_session.assert_called_once()
assert result == (True, {})
@pytest.mark.asyncio
async def test_connect_without_security(opc_repository, mock_client):
def test_connect_without_security(opc_repository, mock_client):
opc_repository.cert_path = None
opc_repository.try_connect = AsyncMock(return_value=(True, {}))
opc_repository.set_security = AsyncMock()
result = await opc_repository.connect()
opc_repository._create_client = MagicMock()
opc_repository._open_session = MagicMock(return_value=(True, {}))
opc_repository.set_security = MagicMock()
result = opc_repository.connect()
opc_repository.try_connect.assert_called_once()
opc_repository._create_client.assert_called_once()
opc_repository._open_session.assert_called_once()
opc_repository.set_security.assert_not_called()
assert opc_repository.client == mock_client
assert result == (True, {})
@pytest.mark.asyncio
async def test_try_connect_success(opc_repository):
opc_repository.last_reconnection_time = None
opc_repository.client = AsyncMock()
result = await opc_repository.try_connect()
def test_connect_raises_when_session_already_open(opc_repository, mock_client):
opc_repository.client = mock_client
proto = MagicMock()
proto.state = 'open'
mock_client.aio_obj.uaclient.protocol = proto
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
opc_repository.connect()
def test_create_client_raises_when_client_exists(opc_repository, mock_client):
opc_repository.client = mock_client
with pytest.raises(OpcClientAlreadyExistsError, match='already exists'):
opc_repository._create_client()
def test_open_session_success(opc_repository):
closed_proto = MagicMock()
closed_proto.state = 'closed'
opc_repository.client = MagicMock()
aio = MagicMock()
opc_repository.client.aio_obj = aio
aio.uaclient = MagicMock(protocol=closed_proto)
aio.session_timeout = 600_000
aio.secure_channel_timeout = 600_000
open_proto = MagicMock()
open_proto.state = 'open'
open_proto.authentication_token = 'tok'
def connect_side_effect():
aio.uaclient.protocol = open_proto
opc_repository.client.connect = MagicMock(side_effect=connect_side_effect)
result = opc_repository._open_session()
opc_repository.client.connect.assert_called_once()
assert opc_repository.last_reconnection_time is not None
assert result == (True, {})
assert opc_repository._session_ready.is_set()
@pytest.mark.asyncio
async def test_try_connect_fail(opc_repository):
opc_repository.last_reconnection_time = None
def test_open_session_raises_when_already_connected(opc_repository, mock_client):
opc_repository.client = mock_client
proto = MagicMock()
proto.state = 'open'
mock_client.aio_obj.uaclient.protocol = proto
with pytest.raises(OpcSessionAlreadyConnectedError, match='disconnect'):
opc_repository._open_session()
def test_open_session_fail(opc_repository):
opc_repository._disconnect_locked = MagicMock()
opc_repository.client = MagicMock()
opc_repository.client.connect.side_effect = Exception("Test error")
aio = MagicMock()
opc_repository.client.aio_obj = aio
aio.uaclient = MagicMock(protocol=MagicMock(state='closed'))
opc_repository.client.connect = MagicMock(side_effect=Exception('Test error'))
is_connected, error_data = await opc_repository.try_connect()
is_connected, error_data = opc_repository._open_session()
opc_repository._disconnect_locked.assert_called_once()
opc_repository.client.connect.assert_called_once()
assert is_connected is False
assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to connect to OPC server: Test error"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
assert error_data['message'] == 'Failed to connect to OPC server: Test error'
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None
@pytest.mark.asyncio
async def test_disconnect(opc_repository, mock_client):
def test_open_session_raises_when_no_client(opc_repository):
opc_repository.client = None
with pytest.raises(OpcClientNotInitializedError, match='not initialized'):
opc_repository._open_session()
def test_disconnection_fallback_success(opc_repository, mock_client):
opc_repository.client = mock_client
await opc_repository.disconnect()
mock_client.disconnect.return_value = True
result = opc_repository._disconnection_fallback()
mock_client.disconnect.assert_called_once()
assert opc_repository.client is None
assert result == []
@pytest.mark.asyncio
async def test_disconnect_no_client(opc_repository):
opc_repository.client = None
assert await opc_repository.disconnect() is None
@pytest.mark.asyncio
async def test_disconnect_error(opc_repository, mock_client):
def test_disconnection_fallback_fail(opc_repository, mock_client):
opc_repository.client = mock_client
mock_client.disconnect.side_effect = Exception("Test error")
await opc_repository.disconnect()
mock_client.disconnect.side_effect = Exception('Test error')
result = opc_repository._disconnection_fallback()
assert result == [
{'attempt': 1, 'error': 'Test error', 'traceback': ANY},
{'attempt': 2, 'error': 'Test error', 'traceback': ANY},
{'attempt': 3, 'error': 'Test error', 'traceback': ANY},
{'attempt': 4, 'error': 'Test error', 'traceback': ANY},
{'attempt': 5, 'error': 'Test error', 'traceback': ANY},
]
assert mock_client.disconnect.call_count == 5
opc_repository.logger.custom_error.assert_called_once_with(
"Failed to disconnect from OPC server: Test error",
ANY
def test_disconnect(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository._disconnection_fallback = MagicMock(return_value=[])
opc_repository.disconnect()
opc_repository._disconnection_fallback.assert_called_once()
assert opc_repository.client is None
assert opc_repository._allow_reconnect is False
def test_disconnect_no_client(opc_repository):
opc_repository.client = None
assert opc_repository.disconnect() is None
def test_disconnect_error(opc_repository, mock_client):
opc_repository.client = mock_client
opc_repository._disconnection_fallback = MagicMock(
return_value=[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}]
)
opc_repository.disconnect()
opc_repository._disconnection_fallback.assert_called_once()
opc_repository.send_notification.assert_called_once_with(
metadata=opc_repository.metadata,
notification_id=f'OPC_DISCONNECTION_ERROR_{opc_repository.id}',
message='Failed to disconnect from OPC server in 5 attempts.',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=json.dumps(
[{'attempt': 1, 'error': 'Test error', 'traceback': 'text'}], indent=4
),
)
assert opc_repository.client is None
@pytest.mark.asyncio
async def test_validate_connection_none_client(opc_repository):
def test_validate_connection_none_client(opc_repository):
opc_repository.client = None
opc_repository.connect = AsyncMock(return_value=(True, {}))
response = await opc_repository.validate_connection()
assert response == (True, {})
opc_repository.connect.assert_called_once()
response = opc_repository.validate_connection()
assert response == (False, opc_repository._not_connected_error())
@pytest.mark.asyncio
async def test_validate_connection_error_count_disconnect_error(opc_repository):
opc_repository.error_count = 6
opc_repository.client = AsyncMock()
opc_repository.disconnect = AsyncMock(
side_effect=Exception("Test error")
)
opc_repository.connect = AsyncMock(return_value=(True, {}))
response = await opc_repository.validate_connection()
assert response == opc_repository.connect.return_value
opc_repository.disconnect.assert_called_once()
opc_repository.connect.assert_called_once()
opc_repository.logger.custom_error.assert_has_calls(
[
call("Failed to disconnect from OPC server: Test error", ANY),
]
)
@pytest.mark.asyncio
async def test_validate_connection_error_validate_connection_error(opc_repository):
opc_repository.client = MagicMock(
uaclient=Exception("Test error")
)
opc_repository.error_count = 0
response = await opc_repository.validate_connection()
assert response == (False, {
"notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}",
"message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
"block": "opc_repository",
"level": NotificationLevel.ERROR,
"attachment_content": ANY
})
@pytest.mark.asyncio
@patch('laborious.utils.repository.opc_repository.datetime')
async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository):
_mock_datetime.now = MagicMock(
return_value=datetime(2025, 1, 1, 0, 0, 0))
opc_repository.error_count = 0
def test_validate_connection_session_not_open(opc_repository):
opc_repository.client = MagicMock()
opc_repository.client.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.connect = MagicMock(return_value=(True, {}))
opc_repository.client.aio_obj.uaclient.protocol = None
response = await opc_repository.validate_connection()
opc_repository.connect.assert_not_called()
assert response == (False, {
"notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}",
"message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...",
"block": "opc_repository",
"level": NotificationLevel.WARNING
})
response = opc_repository.validate_connection()
assert response == (False, opc_repository._not_connected_error())
opc_repository.error.assert_called_once()
@pytest.mark.asyncio
@patch('laborious.utils.repository.opc_repository.datetime')
async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository):
mock_datetime.now = MagicMock(
return_value=datetime(2025, 1, 1, 1, 0, 0))
opc_repository.error_count = 0
opc_repository.client = AsyncMock()
opc_repository.client.uaclient.protocol = None
opc_repository.last_reconnection_time = datetime(2025, 1, 1, 0, 0, 0)
opc_repository.connect = AsyncMock(return_value=(True, {}))
response = await opc_repository.validate_connection()
opc_repository.connect.assert_called_once()
assert response == opc_repository.connect.return_value
@pytest.mark.asyncio
async def test_validate_connection_success(opc_repository):
def test_validate_connection_success(opc_repository):
opc_repository.client = MagicMock()
opc_repository.error_count = 0
opc_repository.client.uaclient.protocol = MagicMock()
opc_repository.client.uaclient.protocol.state = "open"
proto = MagicMock()
proto.state = 'open'
opc_repository.client.aio_obj.uaclient.protocol = proto
output = await opc_repository.validate_connection()
output = opc_repository.validate_connection()
assert output == (True, {})
@pytest.mark.asyncio
async def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
opc_repository.client = AsyncMock(
get_node=MagicMock()
)
mock_node = AsyncMock()
def test_write_data_validate_connection_do_nothing(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = MagicMock(get_node=MagicMock())
mock_node = MagicMock()
opc_repository.client.get_node.return_value = mock_node
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
result = opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
assert result == (True, {})
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert result == (True, {'response_time': ANY})
@pytest.mark.asyncio
async def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = AsyncMock(return_value=(False, {}))
opc_repository.client = AsyncMock()
opc_repository.error_count = 0
def test_write_data_validate_connection_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=(False, {}))
opc_repository.client = MagicMock()
opc_repository._start_reconnect = MagicMock()
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_not_called()
assert result == (False, {})
opc_repository._start_reconnect.assert_called_once()
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
assert is_success is False
assert error_data['opc_error_kind'] == 'connection_lost'
assert error_data['opc_status'] == 'ProtocolClosed'
@pytest.mark.asyncio
async def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
opc_repository.client = AsyncMock()
opc_repository.error_count = 0
opc_repository.client.get_node = MagicMock(
side_effect=Exception("Test error"))
def test_write_data_get_node_failed(opc_repository):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = MagicMock()
opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None
@pytest.mark.asyncio
async def test_write_data_invalid_data_type(opc_repository, mock_client):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
def test_write_data_invalid_data_type(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = mock_client
mock_node = AsyncMock()
mock_node = MagicMock()
mock_client.get_node = MagicMock(return_value=mock_node)
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"invalid_type", opc_repository.logger, metadata['metadata'])
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'invalid_type', metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}"
assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data.get('attachment_content') is None
@pytest.mark.asyncio
@patch('laborious.utils.repository.opc_repository.metrics')
async def test_write_data(mock_metrics, opc_repository, mock_client):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
def test_write_data(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = mock_client
mock_node = AsyncMock()
mock_node = MagicMock()
mock_client.get_node = MagicMock(return_value=mock_node)
result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
result = opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
mock_node.write_value.assert_called_once()
assert result == (True, {})
mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with(
pod_id=opc_repository.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
opc_server_id=opc_repository.id
)
mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
pod_id=opc_repository.pod_id,
model_name=metadata['metadata']['model_name'],
pipeline_name=metadata['metadata']['workflow_name'],
opc_server_id=opc_repository.id
)
mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
ANY)
assert result == (True, {'response_time': ANY})
@pytest.mark.asyncio
async def test_write_data_write_value_failed(opc_repository, mock_client):
opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
def test_write_data_write_value_failed(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = mock_client
mock_node = AsyncMock()
opc_repository.error_count = 0
mock_node = MagicMock()
mock_client.get_node = MagicMock(return_value=mock_node)
mock_node.write_value.side_effect = Exception("Test error")
mock_node.write_value.side_effect = Exception('Test error')
is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
"float", opc_repository.logger, metadata['metadata'])
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
opc_repository.validate_connection.assert_called_once()
mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
mock_node.write_value.assert_called_once()
assert is_success is False
assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}"
assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
assert error_data['block'] == "opc_repository"
assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
assert (
error_data['message']
== "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
)
assert error_data['block'] == 'opc_repository'
assert error_data['level'] == NotificationLevel.ERROR
assert error_data['attachment_content'] is not None
def test_is_reconnectable_opcua_bad():
assert is_reconnectable_opcua_bad(BadSessionIdInvalid()) is True
assert is_reconnectable_opcua_bad(BadNodeIdUnknown()) is False
assert is_reconnectable_opcua_bad(Exception('other')) is False
def test_write_data_bad_session_id_invalid_schedules_reconnect(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = mock_client
opc_repository._start_reconnect = MagicMock()
mock_node = MagicMock()
mock_client.get_node = MagicMock(return_value=mock_node)
mock_node.write_value.side_effect = BadSessionIdInvalid()
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
mock_node.write_value.assert_called_once()
opc_repository._start_reconnect.assert_called_once()
assert is_success is False
assert error_data['opc_error_kind'] == 'session_bad'
assert error_data['opc_status'] == 'BadSessionIdInvalid'
def test_write_data_reconnect_in_progress_immediate(opc_repository):
opc_repository._session_ready.clear()
opc_repository._reconnect_thread = MagicMock()
opc_repository._reconnect_thread.is_alive.return_value = True
opc_repository.validate_connection = MagicMock()
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
opc_repository.validate_connection.assert_not_called()
assert is_success is False
assert error_data['opc_error_kind'] == 'reconnect_in_progress'
def test_start_reconnect_skips_within_interval(opc_repository):
opc_repository.last_reconnection_time = datetime.now()
opc_repository.reconnection_interval = 3600
opc_repository._start_reconnect('BadSessionIdInvalid', 'tok')
assert opc_repository._reconnect_thread is None
def test_write_data_protocol_closed_schedules_reconnect(opc_repository):
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
opc_repository._start_reconnect = MagicMock()
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
opc_repository._start_reconnect.assert_called_once()
assert opc_repository._start_reconnect.call_args.args[0] == 'ProtocolClosed'
assert is_success is False
assert error_data['opc_error_kind'] == 'connection_lost'
assert error_data['opc_status'] == 'ProtocolClosed'
def test_write_data_protocol_closed_skips_reconnect_within_interval(opc_repository):
opc_repository.client = MagicMock()
opc_repository.client.aio_obj.uaclient.protocol = MagicMock(state='closed')
opc_repository.last_reconnection_time = datetime.now()
opc_repository.reconnection_interval = 3600
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
assert opc_repository._reconnect_thread is None
assert is_success is False
assert error_data['opc_error_kind'] == 'connection_lost'
def test_write_data_after_failed_reconnect_schedules_again(opc_repository):
opc_repository._session_ready.clear()
opc_repository.reconnection_interval = 0
opc_repository.last_reconnection_time = None
opc_repository._reconnect_locked = MagicMock(
return_value=(False, {'message': 'connect failed'})
)
opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
if opc_repository._reconnect_thread is not None:
opc_repository._reconnect_thread.join(timeout=2)
assert opc_repository._reconnect_locked.call_count == 1
assert not opc_repository._reconnect_thread_in_progress()
opc_repository.write_data('ns=2;s=TestNode', 42.0, 'float', metadata['metadata'])
if opc_repository._reconnect_thread is not None:
opc_repository._reconnect_thread.join(timeout=2)
assert opc_repository._reconnect_locked.call_count == 2
def test_write_data_after_disconnect_does_not_schedule_reconnect(opc_repository, mock_client):
opc_repository.client = mock_client
proto = MagicMock()
proto.state = 'closed'
mock_client.aio_obj.uaclient.protocol = proto
opc_repository._disconnection_fallback = MagicMock(return_value=[])
opc_repository.disconnect()
is_success, error_data = opc_repository.write_data(
'ns=2;s=TestNode', 42.0, 'float', metadata['metadata']
)
assert opc_repository._reconnect_thread is None
assert is_success is False
assert error_data['opc_error_kind'] == 'connection_lost'
def test_parallel_bad_writes_single_reconnect_task(opc_repository, mock_client):
opc_repository.validate_connection = MagicMock(return_value=(True, {}))
opc_repository.client = mock_client
opc_repository.reconnection_interval = 0
opc_repository.last_reconnection_time = None
mock_node = MagicMock()
mock_client.get_node = MagicMock(return_value=mock_node)
mock_node.write_value.side_effect = BadSessionIdInvalid()
opc_repository._start_reconnect = MagicMock()
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
futures = [
executor.submit(
opc_repository.write_data,
node,
value,
'float',
metadata['metadata'],
)
for node, value in (('ns=2;s=TestNode', 1.0), ('ns=2;s=TestNode2', 2.0))
]
results = [future.result() for future in futures]
assert 1 <= opc_repository._start_reconnect.call_count <= 2
assert mock_node.write_value.call_count == 2
error_kinds = [r[1].get('opc_error_kind') for r in results]
assert error_kinds.count('session_bad') >= 1
assert all(k in ('session_bad', 'reconnect_in_progress') for k in error_kinds)
@patch('laborious.utils.repository.opc_repository.datetime')
def test_reconnect_locked_sets_last_reconnection_time(mock_datetime, opc_repository):
mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 12, 0, 0))
opc_repository._disconnect_locked = MagicMock()
opc_repository._connect_locked = MagicMock(return_value=(True, {}))
result = opc_repository._reconnect_locked()
opc_repository._disconnect_locked.assert_called_once()
opc_repository._connect_locked.assert_called_once()
assert result == (True, {})
assert opc_repository.last_reconnection_time == datetime(2025, 1, 1, 12, 0, 0)

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

@@ -1,14 +1,16 @@
from os import environ
from laborious.utils.connectors_config import (build_mlflow_config,
build_opc_config,
build_postgres_config,
build_mongodb_config)
from laborious.utils.connectors_config import (
build_minio_config,
build_mlflow_config,
build_opc_config,
build_plugin_store_config,
)
def test_build_mlflow_config_with_env_vars():
# Arrange
environ['MLFLOW_HOST'] = 'http://test-host'
environ['MLFLOW_PORT'] = '8080'
environ['MLFLOW_URL'] = 'http://test-host:8080'
environ['MLFLOW_USERNAME'] = 'test-user'
environ['MLFLOW_PASSWORD'] = 'test-pass'
@@ -16,17 +18,25 @@ def test_build_mlflow_config_with_env_vars():
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://test-host'
assert config['port'] == 8080
assert config['url'] == 'http://test-host:8080'
assert config['username'] == 'test-user'
assert config['password'] == 'test-pass'
def test_build_mlflow_config_host_already_has_port():
environ['MLFLOW_URL'] = 'http://tracker.example.com:443'
environ['MLFLOW_USERNAME'] = 'u'
environ['MLFLOW_PASSWORD'] = 'p'
config = build_mlflow_config()
assert config['url'] == 'http://tracker.example.com:443'
def test_build_mlflow_config_with_defaults():
# Arrange
# Clear any existing env vars
environ.pop('MLFLOW_HOST', None)
environ.pop('MLFLOW_PORT', None)
environ.pop('MLFLOW_URL', None)
environ.pop('MLFLOW_USERNAME', None)
environ.pop('MLFLOW_PASSWORD', None)
@@ -34,12 +44,30 @@ def test_build_mlflow_config_with_defaults():
config = build_mlflow_config()
# Assert
assert config['host'] == 'http://localhost'
assert config['port'] == 5080
assert config['url'] == 'http://localhost:5080'
assert config['username'] == 'aignosi'
assert config['password'] == 'aignosi'
def test_build_plugin_store_config_defaults():
environ.pop('STORE_BASE_URL', None)
environ.pop('STORE_OWNER', None)
environ.pop('STORE_REPO', None)
environ.pop('STORE_BRANCH', None)
environ.pop('STORE_USERNAME', None)
environ.pop('STORE_PASSWORD', None)
environ.pop('STORE_CACHE_TTL_SECONDS', None)
environ.pop('PYPI_SERVER', None)
environ.pop('PYPI_USERNAME', None)
environ.pop('PYPI_PASSWORD', None)
cfg = build_plugin_store_config()
assert cfg['base_url'] == 'http://localhost:3000'
assert cfg['owner'] == 'sientia'
assert cfg['repo'] == 'model-library-store'
assert cfg['pypi_index_url'] == 'http://localhost:5000'
def test_build_opc_config_with_env_vars():
# Arrange
environ['OPC_CONFIG'] = '{"opc": {"name": "test-opc", "url": "opc.tcp://test:4840"}}'
@@ -88,74 +116,37 @@ def test_build_opc_config_with_defaults():
assert config['1']['reconnection_interval'] == 120
def test_build_postgres_config_with_env_vars():
# Arrange
environ['POSTGRES_HOST'] = 'test-host'
environ['POSTGRES_PORT'] = '5433'
environ['POSTGRES_USER'] = 'test-user'
environ['POSTGRES_PASSWORD'] = 'test-pass'
environ['POSTGRES_DBNAME'] = 'test-db'
environ['POSTGRES_MIN_CONNECTIONS'] = '10'
environ['POSTGRES_MAX_CONNECTIONS'] = '30'
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'test-host'
assert config['port'] == 5433
assert config['user'] == 'test-user'
assert config['password'] == 'test-pass'
assert config['dbname'] == 'test-db'
assert config['min_connections'] == 10
assert config['max_connections'] == 30
def test_build_postgres_config_with_defaults():
# Arrange
environ.pop('POSTGRES_HOST', None)
environ.pop('POSTGRES_PORT', None)
environ.pop('POSTGRES_USER', None)
environ.pop('POSTGRES_PASSWORD', None)
environ.pop('POSTGRES_DBNAME', None)
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
# Act
config = build_postgres_config()
# Assert
assert config['host'] == 'localhost'
assert config['port'] == 5432
assert config['user'] == 'sientia'
assert config['password'] == 'sientia'
assert config['dbname'] == 'sientia'
assert config['min_connections'] == 5
assert config['max_connections'] == 20
def test_build_mongo_db_config_with_env_vars():
environ['MONGODB_USERNAME'] = 'sientia1'
environ['MONGODB_PASSWORD'] = 'sientia1'
environ['MONGODB_URL'] = 'localhost:27018'
environ['MONGODB_DATABASE_NAME'] = 'test_db'
environ['MONGODB_TTL_INDEX_HOURS'] = '1'
assert build_mongodb_config() == {
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
'database_name': 'test_db',
'ttl_index_seconds': 3600
def test_build_minio_config_with_env_vars():
environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
environ['MINIO_ACCESS_KEY'] = 'test-key'
environ['MINIO_SECRET_KEY'] = 'test-secret'
environ['MINIO_REGION_NAME'] = 'test-region'
environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
# Isolate from IDE/CI env (e.g. VS Code may export MINIO_SECURE=true).
environ['MINIO_SECURE'] = 'false'
assert build_minio_config() == {
'endpoint_url': 'http://test-host',
'access_key': 'test-key',
'secret_key': 'test-secret',
'default_bucket': 'test-bucket',
'retention_hours': 24,
'secure': False,
}
def test_build_mongo_db_config_with_defaults():
environ.pop('MONGODB_USERNAME', None)
environ.pop('MONGODB_PASSWORD', None)
environ.pop('MONGODB_DATABASE_NAME', None)
environ.pop('MONGODB_URL', None)
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
assert build_mongodb_config() == {
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
'database_name': 'sientia',
'ttl_index_seconds': 3600
def test_build_minio_config_with_defaults():
environ.pop('MINIO_ENDPOINT_URL', None)
environ.pop('MINIO_ACCESS_KEY', None)
environ.pop('MINIO_SECRET_KEY', None)
environ.pop('MINIO_REGION_NAME', None)
environ.pop('MINIO_DEFAULT_BUCKET', None)
environ.pop('MINIO_SECURE', None)
environ.pop('MINIO_RETENTION_HOURS', None)
assert build_minio_config() == {
'endpoint_url': 'http://localhost:9000',
'access_key': 'minioadmin',
'secret_key': 'minioadmin',
'default_bucket': 'laborious',
'retention_hours': 24,
'secure': False,
}

12
tests/laborious/utils/test_dataframe_debug.py Normal file
View File

@@ -0,0 +1,12 @@
from pandas import DataFrame
from laborious.utils.dataframe_debug import build_dataframe_debug_message
def test_build_dataframe_debug_message_skips_large_dataframe():
df = DataFrame({'a': [1, 2, 3]})
msg = build_dataframe_debug_message('payload', df, max_rows=1)
assert 'skipped because dataframe has 3 rows' in msg
assert '(max: 1)' in msg

243
tests/laborious/worker/test_worker.py Normal file
View File

@@ -0,0 +1,243 @@
from unittest.mock import AsyncMock, MagicMock, patch
from pytest import mark, raises
from laborious.worker import worker
def _build_fake_activities():
inst = MagicMock()
inst.init_opc = MagicMock()
inst.shutdown = MagicMock()
inst.load_query_with_minio_offload = MagicMock()
inst.retrain_model = MagicMock()
inst.update_production_model = MagicMock()
inst.format_retrain_report = MagicMock()
inst.export_data_to_postgres = MagicMock()
inst.load_custom_query = MagicMock()
inst.calculate_simple_metrics = MagicMock()
inst.get_reference_data = MagicMock()
inst.calculate_drift = MagicMock()
inst.request_predict = MagicMock()
inst.request_transform = MagicMock()
inst.input_gate = MagicMock()
inst.mlflow_response_gate = MagicMock()
inst.mlflow_content_gate = MagicMock()
inst.format_transformed_data = MagicMock()
inst.format_prediction = MagicMock()
inst.format_default_prediction = MagicMock()
inst.write_opc_data = MagicMock()
inst.cleanup_minio_objects_expired = MagicMock()
inst.repeat_last_prediction = MagicMock()
inst.write_metrics = MagicMock()
inst.write_pi_web_api_data = MagicMock()
return inst
def _build_fake_worker(async_result=None, async_error: Exception | None = None):
w = MagicMock()
async def _run():
if async_error is not None:
raise async_error
return async_result
w.run = MagicMock(side_effect=_run)
return w
@patch('laborious.worker.worker.start_http_server')
def test_start_prometheus_server_success(mock_start_http):
with patch.object(worker.metrics.APP_UP, 'labels') as labels:
gauge = MagicMock()
labels.return_value = gauge
with patch('laborious.worker.worker.os.getenv', return_value='9090'):
worker.start_prometheus_server()
mock_start_http.assert_called_once_with(9090)
gauge.set.assert_called_once_with(1)
@patch('laborious.worker.worker.start_http_server', side_effect=RuntimeError('nope'))
def test_start_prometheus_server_error_exits(_mock_start_http):
with patch('laborious.worker.worker.os._exit', side_effect=SystemExit(1)) as m_exit:
with raises(SystemExit):
worker.start_prometheus_server()
m_exit.assert_called_once_with(1)
@mark.asyncio
async def test_main_missing_runtime_exits_fast(monkeypatch):
monkeypatch.setenv('RUNTIME', '')
with (
patch('laborious.worker.worker.start_prometheus_server'),
patch('laborious.worker.worker.get_logger') as m_logger,
patch('laborious.worker.worker.NotificationHandler'),
patch('laborious.worker.worker.MetricsController'),
patch.object(worker.metrics.APP_UP, 'labels') as labels,
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
):
labels.return_value = MagicMock()
with raises(SystemExit):
await worker.main()
assert m_logger.return_value.custom_critical.called
@mark.asyncio
async def test_main_plugin_install_failure(monkeypatch):
monkeypatch.setenv('RUNTIME', 'single')
fake_activities = _build_fake_activities()
fake_plugin = MagicMock()
fake_plugin.install_runtime = AsyncMock(side_effect=RuntimeError('install failed'))
with (
patch('laborious.worker.worker.start_prometheus_server'),
patch('laborious.worker.worker.get_logger'),
patch(
'laborious.worker.worker.build_mongodb_config',
return_value={'connection_string': 'cs', 'database_name': 'db'},
),
patch('laborious.worker.worker.NotificationHandler'),
patch('laborious.worker.worker.MetricsController'),
patch(
'laborious.worker.worker.build_plugin_store_config',
return_value={
'base_url': '',
'owner': '',
'repo': '',
'username': None,
'password': None,
'branch': None,
'cache_ttl_seconds': None,
'pypi_index_url': '',
'pypi_username': None,
'pypi_password': None,
},
),
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
patch('laborious.worker.worker.Activities', return_value=fake_activities),
patch.object(worker.metrics.APP_UP, 'labels') as labels,
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
):
labels.return_value = MagicMock()
with raises(SystemExit):
await worker.main()
@mark.asyncio
async def test_main_success_exit_zero(monkeypatch):
monkeypatch.setenv('RUNTIME', 'single')
fake_activities = _build_fake_activities()
fake_plugin = MagicMock()
fake_plugin.install_runtime = AsyncMock(return_value=None)
fake_workers = [_build_fake_worker() for _ in range(4)]
with (
patch('laborious.worker.worker.start_prometheus_server'),
patch('laborious.worker.worker.get_logger'),
patch(
'laborious.worker.worker.build_mongodb_config',
return_value={'connection_string': 'cs', 'database_name': 'db'},
),
patch('laborious.worker.worker.NotificationHandler') as m_notif_cls,
patch('laborious.worker.worker.MetricsController'),
patch(
'laborious.worker.worker.build_plugin_store_config',
return_value={
'base_url': '',
'owner': '',
'repo': '',
'username': None,
'password': None,
'branch': None,
'cache_ttl_seconds': None,
'pypi_index_url': '',
'pypi_username': None,
'pypi_password': None,
},
),
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
patch('laborious.worker.worker.Activities', return_value=fake_activities),
patch('laborious.worker.worker.build_postgres_config', return_value={}),
patch('laborious.worker.worker.build_minio_config', return_value={}),
patch('laborious.worker.worker.build_opc_config', return_value={}),
patch('laborious.worker.worker.build_api_config', return_value={}),
patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
patch(
'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
),
patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers) as m_prepare,
patch.object(worker.metrics.APP_UP, 'labels') as labels,
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(0)),
):
labels.return_value = MagicMock()
with raises(SystemExit):
await worker.main()
notif = m_notif_cls.return_value
notif.shutdown.assert_called_once()
fake_activities.shutdown.assert_called_once()
assert m_prepare.call_count == 4
prepare_calls = m_prepare.call_args_list
assert prepare_calls[0].kwargs['runtime'] == 'single'
assert prepare_calls[3].kwargs['runtime'] == 'single'
@mark.asyncio
async def test_main_worker_gather_error_exits_one(monkeypatch):
monkeypatch.setenv('RUNTIME', 'single')
fake_activities = _build_fake_activities()
fake_plugin = MagicMock()
fake_plugin.install_runtime = AsyncMock(return_value=None)
fake_workers = [
_build_fake_worker(async_error=RuntimeError('boom')),
_build_fake_worker(),
_build_fake_worker(),
_build_fake_worker(),
]
with (
patch('laborious.worker.worker.start_prometheus_server'),
patch('laborious.worker.worker.get_logger') as m_logger,
patch(
'laborious.worker.worker.build_mongodb_config',
return_value={'connection_string': 'cs', 'database_name': 'db'},
),
patch('laborious.worker.worker.NotificationHandler'),
patch('laborious.worker.worker.MetricsController'),
patch(
'laborious.worker.worker.build_plugin_store_config',
return_value={
'base_url': '',
'owner': '',
'repo': '',
'username': None,
'password': None,
'branch': None,
'cache_ttl_seconds': None,
'pypi_index_url': '',
'pypi_username': None,
'pypi_password': None,
},
),
patch('laborious.worker.worker.PluginStore', return_value=fake_plugin),
patch('laborious.worker.worker.Activities', return_value=fake_activities),
patch('laborious.worker.worker.build_postgres_config', return_value={}),
patch('laborious.worker.worker.build_minio_config', return_value={}),
patch('laborious.worker.worker.build_opc_config', return_value={}),
patch('laborious.worker.worker.build_api_config', return_value={}),
patch('laborious.worker.worker.PrometheusConfig', return_value=MagicMock()),
patch('laborious.worker.worker.TelemetryConfig', return_value=MagicMock()),
patch('laborious.worker.worker.Runtime', return_value=MagicMock()),
patch(
'laborious.worker.worker.client.Client.connect', new=AsyncMock(return_value=MagicMock())
),
patch('laborious.worker.worker.prepare_worker', side_effect=fake_workers),
patch.object(worker.metrics.APP_UP, 'labels') as labels,
patch('laborious.worker.worker.sys.exit', side_effect=SystemExit(1)),
):
labels.return_value = MagicMock()
with raises(SystemExit):
await worker.main()
assert m_logger.return_value.custom_error.called

755
tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py
View File

@@ -1,9 +1,10 @@
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -12,149 +13,667 @@ def format_and_export_prediction():
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
"path_flag": None,
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"prediction_store_policy": "erl:1"
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_servers': ['test_server'],
'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'erl:1',
}
prediction_data = MagicMock()
opc_metrics = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
(prediction_data, opc_metrics),
MagicMock(),
MagicMock(),
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_prediction,
{
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
**metadata,
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value,
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': prediction_data,
'opc_metrics': opc_metrics,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag_with_transformed_data(
workflow_mock, format_and_export_prediction
):
# Arrange
input_data = {
'metadata': metadata,
"path_flag": "default",
"data": {"test": "data"},
"timestamp": "2021-01-01",
"model_id": 1,
"prediction_confidence": 0,
"schema": "test_schema",
"table_name": "test_table",
"opc_servers": ["test_server"],
"opc_output_config": {"test": "config"},
"comment": "test_comment"
'path_flag': None,
'data': {'test': 'data'},
'transformed_data': {'transformed': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0.9,
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'opc_servers': ['test_server'],
'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'lts:1',
}
prediction_data = MagicMock()
opc_metrics = MagicMock()
transformed_data = MagicMock()
workflow_mock.execute_local_activity_method.side_effect = [
prediction_data, # format_prediction
transformed_data, # format_transformed_data
]
write_transformed_handler = AsyncMock()
workflow_mock.start_activity_method.return_value = write_transformed_handler
workflow_mock.execute_activity_method.side_effect = [
(prediction_data, opc_metrics), # write_opc_data
MagicMock(), # export_data_to_postgres (prediction)
MagicMock(), # write_metrics
]
# Act
await format_and_export_prediction.run(input_data)
# Assert - format_prediction call
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
**metadata,
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.format_transformed_data,
{
**metadata,
'data': input_data['transformed_data'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
# Assert - start_activity_method for transformed data export
workflow_mock.start_activity_method.assert_called_once_with(
Activities.export_payload_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['transform_table_name'],
'data': transformed_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
# Assert - write_opc_data call
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': prediction_data,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
# Assert - export_data_to_postgres for prediction call
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
# Assert - write_metrics call
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': prediction_data,
'opc_metrics': opc_metrics,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
# Assert - verify counts
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 2
assert workflow_mock.start_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_servers': ['test_server'],
'opc_output_config': {'test': 'config'},
'comment': 'test_comment',
}
prediction_data = MagicMock()
opc_metrics = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
(prediction_data, opc_metrics),
MagicMock(),
MagicMock(),
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_default_prediction,
{
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_default_prediction,
{
**metadata,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value,
**metadata,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ
}
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': prediction_data,
'opc_metrics': opc_metrics,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com',
'prediction_tags': {},
'confidence_tags': {},
},
'prediction_store_policy': 'erl:1',
}
pi_web_api_data = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_prediction,
{
**metadata,
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'prediction_store_policy': input_data['prediction_store_policy'],
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': pi_web_api_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': pi_web_api_data,
'opc_metrics': {},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_none_path_flag_with_pi_web_api_and_opc(
workflow_mock, format_and_export_prediction
):
input_data = {
'metadata': metadata,
'path_flag': None,
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'opc_output_config': {'test': 'config'},
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com',
'prediction_tags': {},
'confidence_tags': {},
},
'prediction_store_policy': 'erl:1',
}
prediction_data = MagicMock()
pi_web_api_data = MagicMock()
opc_metrics = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
(prediction_data, opc_metrics), # write_opc_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': pi_web_api_data,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_metrics,
{
**metadata,
'prediction': prediction_data,
'opc_metrics': opc_metrics,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 4
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@patch(
'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
new_callable=AsyncMock,
)
async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_export_prediction):
input_data = {
'metadata': metadata,
'path_flag': 'default',
'data': {'test': 'data'},
'timestamp': '2021-01-01',
'model_id': 1,
'model_name': metadata['metadata']['model_name'],
'prediction_confidence': 0,
'schema': 'test_schema',
'table_name': 'test_table',
'pi_web_api_output_config': {
'endpoint': 'https://test-pi-server.com',
'prediction_tags': {},
'confidence_tags': {},
},
'comment': 'test_comment',
}
pi_web_api_data = MagicMock()
workflow_mock.execute_activity_method.side_effect = [
pi_web_api_data, # write_pi_web_api_data
MagicMock(), # export_data_to_postgres
MagicMock(), # write_metrics
]
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': pi_web_api_data,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
'on_conflict': 'error',
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

1093
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

File diff suppressed because it is too large Load Diff

248
tests/laborious/workflows/test_drift.py Normal file
View File

@@ -0,0 +1,248 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.drift import Drift
@fixture
def drift() -> Drift:
return Drift()
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'drift',
'schedule_name': 'test_schedule',
},
}
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
'chunk_period': 'hour',
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
# Act
await drift.run(input_data)
# Assert - Check start_local_activity_method calls
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_gathering_query = f"""
SELECT *
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
WHERE
model_id = '{input_data['model_id']}' AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
ORDER BY timestamp ASC
"""
workflow_mock.start_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': expected_gathering_query,
'datetime_columns': ['timestamp', 'created_at'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.get_reference_data,
{
**metadata,
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
# Assert - Check calculate_drift call
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data['chunk_period'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
# Assert - Check export_data_to_postgres call
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.export_data_to_postgres,
{
**metadata,
'data': drift_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
}
target_data = None
reference_data = {'data': 'test_reference_data'}
workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_activity_method = AsyncMock()
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Should not call calculate_drift or export
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = None
workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
workflow_mock.execute_activity_method = AsyncMock()
# Act
await drift.run(input_data)
# Assert - Should call calculate_drift but not export
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data.get('chunk_period', 'min'),
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'schema': 'test_schema',
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
# chunk_period not provided, should default to 'min'
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
# Act
await drift.run(input_data)
# Assert - Check calculate_drift call with default chunk_period
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': 'min', # Default value
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)

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

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.minimal_retrain import MinimalRetrain
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
},
}
@@ -23,76 +25,300 @@ metadata = {
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "minimal_retrain",
"schedule_name": "test_schedule",
"query": "test_query",
"schema": "test_schema",
"table_name": "test_table",
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'retention_minutes': 0,
},
}
storage_result = {
'last_timestamp': '2024-01-01 00:00:00+0000',
'status': {'success': True},
'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
'bucket': None,
'object_key': None,
'object_prefix': None,
'uri': None,
}
workflow_mock.execute_activity_method = AsyncMock(
return_value={
"data1": "1",
"data2": "2",
}
side_effect=[
storage_result,
{'success': True, 'experiment': 'test_experiment'},
{
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
{'report': 'test_report'},
]
)
await minimal_retrain.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls(
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
Activities.load_query_with_minio_offload,
{
**metadata,
"query": input_data["query"],
'datetime_columns': input_data.get('datetime_columns', [])
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.retrain_model,
{
**metadata,
'data': workflow_mock.execute_local_activity_method.return_value,
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.retrain_model,
{
**metadata,
'data': storage_result,
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.update_production_model,
{
**metadata,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
**workflow_mock.execute_activity_method.return_value,
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.update_production_model,
{
**metadata,
'model_name': input_data['model_name'],
'success': True,
'experiment': 'test_experiment',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.export_data_to_postgres,
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': {'success': True, 'experiment': 'test_experiment'},
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'update_report': {
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'data': workflow_mock.execute_local_activity_method.return_value,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
@mark.asyncio
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'retention_minutes': 0,
},
}
storage_result = {
'last_timestamp': '2024-01-01 00:00:00+0000',
'status': {'success': True},
'data': {},
'bucket': None,
'object_key': None,
'object_prefix': None,
'uri': None,
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
storage_result,
{'success': True, 'experiment': 'test_experiment'},
{
**metadata,
'data': workflow_mock.execute_activity_method.return_value,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
{'report': 'test_report'},
]
)
from pytest import raises
with raises(ValueError, match='No data returned from query'):
await minimal_retrain.run(input_data)
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_local_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
input_data = {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'minimal_retrain',
'schedule_name': 'test_schedule',
'query': 'test_query',
'schema': 'test_schema',
'table_name': 'test_table',
'model_config': {
'target': 'test_target',
'retention_minutes': 0,
},
}
storage_result = {
'last_timestamp': '2024-01-01 00:00:00+0000',
'status': {'success': True},
'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}},
'bucket': None,
'object_key': None,
'object_prefix': None,
'uri': None,
}
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
storage_result,
{'success': False, 'experiment': 'test_experiment'},
{
'success': True,
'version': 'test_version',
'mlflow_run_id': 'test_mlflow_run_id',
'mlflow_experiment_id': 'test_mlflow_experiment_id',
},
{'report': 'test_report'},
]
)
await minimal_retrain.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.retrain_model,
{
**metadata,
'data': storage_result,
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': {'success': False, 'experiment': 'test_experiment'},
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'update_report': {},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.export_data_to_postgres,
{
**metadata,
'data': workflow_mock.execute_local_activity_method.return_value,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_local_activity_method.call_count == 1

112
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -1,5 +1,7 @@
from unittest.mock import AsyncMock, call, patch, ANY
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.predictions_batch import PredictionsBatch
@@ -10,21 +12,19 @@ def predictions_batch() -> PredictionsBatch:
metadata = {
"metadata": {
"model_id": "test_model_id",
"model_name": "test_model",
"workflow_name": "predictions_batch",
"schedule_name": "test_schedule",
},
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'predictions_batch',
'schedule_name': 'test_schedule',
}
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
activity_return = MagicMock()
workflow_mock.execute_activity_method.return_value = activity_return
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
@@ -32,57 +32,75 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'query': 'SELECT * FROM test',
'schema': 'test_schema',
'table_name': 'test_table',
'transform_table_name': 'test_transform_table',
'opc_output_config': 'test_opc_output_config',
'pi_web_api_output_config': 'test_pi_web_api_output_config',
'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {
'retention': '30'
}
'model_config': {'retention': '30'},
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', [])
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_query_with_minio_offload,
{
'metadata': metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
]
)
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
'metadata': {'metadata': metadata},
'data': activity_return,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {
'EMPTY_DATA': {
'POLICY': 'STOP'
}
}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
'API_ERROR': {
'POLICY': 'STOP'
}
}),
'input_filters': input_data.get(
'input_filters',
{
'EMPTY_DATA': {
'POLICY': 'STOP',
'CONFIG': {},
}
},
),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters',
{
'API_ERROR': {
'POLICY': 'STOP',
'CONFIG': {},
}
},
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters',
{
'API_ERROR': {
'POLICY': 'STOP',
'CONFIG': {},
}
},
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
'on_conflict': input_data.get('on_conflict', 'error'),
'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
'save_transform': input_data.get('save_transform', True),
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])
workflow_mock.execute_child_workflow.assert_has_calls(
[call('subworkflow.prediction_process', prediction_input)]
)

215
tests/laborious/workflows/test_simple_metrics.py Normal file
View File

@@ -0,0 +1,215 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.simple_metrics import SimpleMetrics
@fixture
def simple_metrics() -> SimpleMetrics:
return SimpleMetrics()
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'simple_metrics',
'schedule_name': 'test_schedule',
},
}
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse', 'mae', 'r2'],
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None])
workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
# Act
await simple_metrics.run(input_data)
# Assert - Check load_custom_query call
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
on p."timestamp" = ld."timestamp"
where
p.model_id = '{input_data['model_id']}' and
p.prediction is not null and
ld.variable = '{input_data['model_config']['target']}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
order by
p."timestamp" desc;
"""
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': expected_query,
'datetime_columns': ['timestamp'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.export_data_to_postgres,
{
**metadata,
'data': simple_metrics_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': input_data['model_id'],
'target_data': target_data,
'metrics': input_data['metrics'],
'interval_minutes': input_data['interval_minutes'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse'],
}
target_data = None
workflow_mock.execute_activity_method = AsyncMock(return_value=target_data)
# Act
await simple_metrics.run(input_data)
# Assert - Should not call calculate_simple_metrics or export
assert workflow_mock.execute_activity_method.call_count == 1
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse'],
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = None
workflow_mock.execute_activity_method = AsyncMock(return_value=target_data)
workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
# Act
await simple_metrics.run(input_data)
# Assert - Should call calculate_simple_metrics but not export
workflow_mock.execute_activity_method.assert_called_once()
workflow_mock.execute_local_activity_method.assert_called_once()
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
# metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2']
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None])
workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
# Act
await simple_metrics.run(input_data)
# Assert - Check calculate_simple_metrics call with default metrics
workflow_mock.execute_activity_method.assert_any_call(
Activities.load_custom_query,
ANY,
retry_policy=ANY,
start_to_close_timeout=ANY,
)
workflow_mock.execute_local_activity_method.assert_called_once_with(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': input_data['model_id'],
'target_data': target_data,
'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value
'interval_minutes': input_data['interval_minutes'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)