Code import - branch release/SIENTIAPDE-1646
This commit is contained in:
@@ -0,0 +1,679 @@
|
||||
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
|
||||
|
||||
|
||||
@fixture
|
||||
def format_and_export_prediction():
|
||||
return FormatAndExportPrediction()
|
||||
|
||||
|
||||
metadata = {
|
||||
'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,
|
||||
)
|
||||
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,
|
||||
'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,
|
||||
{
|
||||
**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.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_transformed_data(
|
||||
workflow_mock, format_and_export_prediction
|
||||
):
|
||||
# Arrange
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'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,
|
||||
{
|
||||
**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.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
|
||||
@@ -0,0 +1,866 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
|
||||
|
||||
@fixture
|
||||
def prediction_process():
|
||||
return PredictionProcess()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
'schedule_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(return_value=False)
|
||||
# Arrange
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {'test': 'config'},
|
||||
'prediction_store_policy': 'lts:1',
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
{'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
# mlflow_response_gate (predict)
|
||||
('continue', 0.95, 'Error'),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 4
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
**metadata,
|
||||
'data': input_data['data'],
|
||||
'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.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_predict,
|
||||
{
|
||||
**metadata,
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'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.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
'filters': input_data['mlflow_predict_filters'],
|
||||
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'predict',
|
||||
'path_priority': input_data['path_priority'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'on_conflict': 'error',
|
||||
'path_flag': 'continue',
|
||||
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'prediction_confidence': 0.95,
|
||||
'timestamp': '2024-01-01',
|
||||
'model_id': 1,
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': input_data['model_config'],
|
||||
'opc_output_config': input_data['opc_output_config'],
|
||||
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'transform_table_name': input_data['transform_table_name'],
|
||||
'comment': 'Error',
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(return_value=True)
|
||||
# Arrange
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('stop', 0.95, 'Input data with bad quality'), # input_gate
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
),
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
|
||||
# Arrange
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('repeat', 0.95, 'Input data with bad quality'), # input_gate
|
||||
('continue', 0.95, 'Error'), # mlflow_response_gate (transform)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 2
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
|
||||
# Arrange
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 3
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
|
||||
# Arrange
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
}
|
||||
|
||||
# Mock the activity responses
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
{'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'Input data with bad quality'), # input_gate
|
||||
# mlflow_response_gate (transform)
|
||||
('continue', 0.95, 'Error'),
|
||||
# mlflow_content_gate (transform)
|
||||
('continue', 0.95, 'Transformed data not passed the content filter'),
|
||||
('continue', 0.95, 'Error'), # mlflow_response_gate (predict)
|
||||
]
|
||||
|
||||
# Act
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
# Assert
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 4
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.input_gate,
|
||||
{
|
||||
'filters': input_data['input_filters'],
|
||||
'data': input_data['data'],
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_transform,
|
||||
{
|
||||
'data': input_data['data'],
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_transform_filters'],
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'transform',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.request_predict,
|
||||
{
|
||||
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
'model_name': input_data['model_name'],
|
||||
'model_config': input_data['model_config'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
'filters': input_data['mlflow_predict_filters'],
|
||||
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
'type': 'predict',
|
||||
'path_priority': input_data['path_priority'],
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'STOP'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'repeat'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'model': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'CONTINUE'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
prediction_store_policy = 'erl:1'
|
||||
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
'metadata': metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com',
|
||||
'prediction_tags': {},
|
||||
'confidence_tags': {},
|
||||
},
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'Prediction Process',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is True
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_called_once_with(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'path_flag': path_flag,
|
||||
'data': data,
|
||||
'prediction_confidence': confidence,
|
||||
'timestamp': last_timestamp,
|
||||
'model_id': model,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'comment': 'Prediction Process',
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com',
|
||||
'prediction_tags': {},
|
||||
'confidence_tags': {},
|
||||
},
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
'on_conflict': 'error',
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
|
||||
# Arrange
|
||||
data = {'test': 'data'}
|
||||
path_flag = 'unknown'
|
||||
confidence = 0.95
|
||||
schema = 'test_schema'
|
||||
table_name = 'test_table'
|
||||
model = 'test_model'
|
||||
last_timestamp = '2024-01-01'
|
||||
model_name = 'test_model_name'
|
||||
model_config = {'retention': '30'}
|
||||
prediction_store_policy = 'erl:1'
|
||||
# Act
|
||||
result = await prediction_process.path_flag_handler(
|
||||
data,
|
||||
path_flag,
|
||||
{
|
||||
**metadata,
|
||||
'schema': schema,
|
||||
'table_name': table_name,
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': model,
|
||||
'last_timestamp': last_timestamp,
|
||||
'model_name': model_name,
|
||||
'model_config': model_config,
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'https://test-pi-server.com',
|
||||
'prediction_tags': {},
|
||||
'confidence_tags': {},
|
||||
},
|
||||
'prediction_store_policy': prediction_store_policy,
|
||||
},
|
||||
confidence,
|
||||
last_timestamp,
|
||||
'',
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert result is False
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
workflow_mock.execute_child_workflow.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process):
|
||||
prediction_process.path_flag_handler = AsyncMock(return_value=False)
|
||||
prediction_process.cleanup_prefixes = {'training_datasets/test'}
|
||||
|
||||
data_payload = MagicMock()
|
||||
data_payload.cleanup_prefix.return_value = 'training_datasets/test'
|
||||
data_payload.__getitem__ = (
|
||||
lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
|
||||
)
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': data_payload,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'transform_table_name': 'test_transform_table',
|
||||
'model_id': 1,
|
||||
'input_filters': {'test': 'filter'},
|
||||
'mlflow_transform_filters': {'test': 'filter'},
|
||||
'mlflow_predict_filters': {'test': 'filter'},
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'retention': '30'},
|
||||
'path_priority': ['continue', 'repeat', 'stop'],
|
||||
'opc_output_config': {'test': 'config'},
|
||||
'pi_web_api_output_config': {'test': 'config'},
|
||||
'prediction_store_policy': 'lts:1',
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method.side_effect = [
|
||||
{'content': 'transformed_data', 'timestamp': '2024-01-01'},
|
||||
{'content': 'predicted_data', 'timestamp': '2024-01-01'},
|
||||
MagicMock(),
|
||||
]
|
||||
workflow_mock.execute_local_activity_method.side_effect = [
|
||||
('continue', 0.95, 'ok'),
|
||||
('continue', 0.95, ''),
|
||||
('continue', 0.95, ''),
|
||||
('continue', 0.95, ''),
|
||||
]
|
||||
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_any_call(
|
||||
Activities.cleanup_minio_objects_expired,
|
||||
{**metadata, 'data': data_payload},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
|
||||
async def test_run_always_cleans_up_on_pipeline_exception(workflow_mock, prediction_process):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'data': {'last_timestamp': '2024-01-01'},
|
||||
'model_id': 1,
|
||||
'model_name': 'm',
|
||||
'model_config': {},
|
||||
'save_transform': False,
|
||||
}
|
||||
prediction_process._run_prediction_pipeline = AsyncMock(side_effect=RuntimeError('boom'))
|
||||
|
||||
with raises(RuntimeError):
|
||||
await prediction_process.run(input_data)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_called_once_with(
|
||||
Activities.cleanup_minio_objects_expired,
|
||||
{**metadata, 'data': input_data['data']},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
248
tests/laborious/workflows/test_drift.py
Normal file
248
tests/laborious/workflows/test_drift.py
Normal 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,
|
||||
)
|
||||
324
tests/laborious/workflows/test_minimal_retrain.py
Normal file
324
tests/laborious/workflows/test_minimal_retrain.py
Normal file
@@ -0,0 +1,324 @@
|
||||
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
|
||||
|
||||
|
||||
@fixture
|
||||
def minimal_retrain() -> MinimalRetrain:
|
||||
return MinimalRetrain()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'workflow_name': 'minimal_retrain',
|
||||
'schedule_name': 'test_schedule',
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@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_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': 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_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_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_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'},
|
||||
{
|
||||
'success': True,
|
||||
'version': 'test_version',
|
||||
'mlflow_run_id': 'test_mlflow_run_id',
|
||||
'mlflow_experiment_id': 'test_mlflow_experiment_id',
|
||||
},
|
||||
{'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
|
||||
106
tests/laborious/workflows/test_predictions_batch.py
Normal file
106
tests/laborious/workflows/test_predictions_batch.py
Normal file
@@ -0,0 +1,106 @@
|
||||
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
|
||||
|
||||
|
||||
@fixture
|
||||
def predictions_batch() -> PredictionsBatch:
|
||||
return PredictionsBatch()
|
||||
|
||||
|
||||
metadata = {
|
||||
'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):
|
||||
activity_return = MagicMock()
|
||||
workflow_mock.execute_activity_method.return_value = activity_return
|
||||
|
||||
input_data = {
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'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'},
|
||||
}
|
||||
|
||||
await predictions_batch.run(input_data)
|
||||
|
||||
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': 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',
|
||||
'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', {}),
|
||||
'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('subworkflow.prediction_process', prediction_input)]
|
||||
)
|
||||
215
tests/laborious/workflows/test_simple_metrics.py
Normal file
215
tests/laborious/workflows/test_simple_metrics.py
Normal 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,
|
||||
)
|
||||
Reference in New Issue
Block a user