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