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

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