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

SIENTIAPDE-1222

Browse Source 
Refactor model configuration handling in MLFlow and workflows

- Replaced 'model_retention' with 'model_config' to encapsulate retention settings and improve consistency across various components.
- Updated test cases to reflect changes in argument structure, ensuring compatibility with the new model configuration format.
- Added 'prediction_store_policy' to input data handling in workflows for enhanced configuration management.
This commit is contained in:
vitor-aignosi
2025-09-17 16:02:54 -03:00
parent d6fdbc58bb
commit 84c1371d6c
4 changed files with 111 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
## Problemas no Courier:.md
View File

File diff suppressed because one or more lines are too long

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

@@ -4,6 +4,7 @@ from unittest.mock import ANY, MagicMock, patch
import numpy as np import numpy as np
from pandas import DataFrame, Timestamp from pandas import DataFrame, Timestamp
from pytest import fixture, mark, raises from pytest import fixture, mark, raises
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious.activities.mlflow import MLFlow from laborious.activities.mlflow import MLFlow
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
@@ -79,7 +80,7 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
'value': 1.0, 'created_at': '2024-01-01 12:00:00'} 'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
], ],
'model_name': 'test_model', 'model_name': 'test_model',
'model_retention': 30 'model_config': {}
} }
# Mock the transform response # Mock the transform response
@@ -108,26 +109,35 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
# Verify the repository was called with correct arguments # Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.transform.assert_called_once_with( mlflow.model_monitoring_repository.transform.assert_called_once_with(
'test_model', mock_dataframe, 30 'test_model', mock_dataframe, {}, metadata['metadata']
) )
@mark.asyncio @mark.asyncio
@patch("laborious.activities.mlflow.DataFrame") @patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.to_datetime")
@patch("laborious.activities.mlflow.max") @patch("laborious.activities.mlflow.max")
async def test_request_predict(mock_max, mock_dataframe, mlflow): async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
mock_max.return_value = '2024-01-02' mock_max.return_value = '2024-01-02'
# Mock input data # Mock input data
input_data = { input_data = {
**metadata, **metadata,
'data': [ 'data': {
{'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0}, "variable": {
{'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0}, "2024-01-01": "var1",
{'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0}, "2024-01-02": "var2",
{'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0} "2024-01-03": "var1",
], "2024-01-04": "var2"
},
"value": {
"2024-01-01": 1.0,
"2024-01-02": 2.0,
"2024-01-03": 3.0,
"2024-01-04": 4.0
}
},
'model_name': 'test_model', 'model_name': 'test_model',
'model_retention': 30 'model_config': {}
} }
# Mock the predict response # Mock the predict response
@@ -141,13 +151,26 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow):
mock_dataframe.return_value.replace.assert_called_once_with( mock_dataframe.return_value.replace.assert_called_once_with(
np.nan, None, inplace=True np.nan, None, inplace=True
) )
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_dataframe.return_value.__setitem__.assert_any_call(
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
)
mock_to_datetime.assert_called_once_with(
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
)
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
DATETIME_FORMAT
)
# Verify the response # Verify the response
assert response_data == expected_response assert response_data == expected_response
# Verify the repository was called with correct arguments # Verify the repository was called with correct arguments
mlflow.model_monitoring_repository.predict.assert_called_once_with( mlflow.model_monitoring_repository.predict.assert_called_once_with(
'test_model', mock_dataframe.return_value, 30 'test_model', mock_dataframe.return_value, {}, metadata['metadata']
) )

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

@@ -34,7 +34,9 @@ async def test_run(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'], 'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}, 'opc_output_config': {'test': 'config'},
'prediction_store_policy': 'lts:1' 'prediction_store_policy': 'lts:1'
@@ -62,54 +64,54 @@ async def test_run(workflow_mock, prediction_process):
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, { call(Activities.get_last_timestamp, {
**metadata,
'data': input_data['data'], 'data': input_data['data'],
**metadata
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, { call(Activities.input_gate, {
**metadata,
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': input_data['data'], 'data': input_data['data'],
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, { call(Activities.request_transform, {
**metadata,
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
**metadata,
'filters': input_data['mlflow_transform_filters'], 'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
**metadata,
'filters': input_data['mlflow_transform_filters'], 'filters': input_data['mlflow_transform_filters'],
'data': 'transformed_data', 'data': 'transformed_data',
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, { call(Activities.request_predict, {
**metadata,
'data': 'transformed_data', 'data': 'transformed_data',
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, { call(Activities.mlflow_response_gate, {
**metadata,
'filters': input_data['mlflow_predict_filters'], 'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict', 'type': 'predict',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with( workflow_mock.execute_child_workflow.assert_called_once_with(
@@ -122,11 +124,12 @@ async def test_run(workflow_mock, prediction_process):
'timestamp': '2024-01-01', 'timestamp': '2024-01-01',
'model_id': 1, 'model_id': 1,
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': input_data['model_config'],
'opc_output_config': input_data['opc_output_config'], 'opc_output_config': input_data['opc_output_config'],
'schema': input_data['schema'], 'schema': input_data['schema'],
'table_name': input_data['table_name'], 'table_name': input_data['table_name'],
'comment': 'Error' 'comment': 'Error',
'prediction_store_policy': input_data['prediction_store_policy']
} }
) )
@@ -146,7 +149,9 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'], 'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'}
} }
@@ -165,13 +170,13 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, { call(Activities.get_last_timestamp, {
'data': input_data['data'], 'data': input_data['data'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY), }, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, { call(Activities.input_gate, {
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': input_data['data'], 'data': input_data['data'],
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY) }, retry_policy=ANY, start_to_close_timeout=ANY)
]) ])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -192,7 +197,9 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'], 'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'}
} }
@@ -213,7 +220,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, { call(Activities.get_last_timestamp, {
'data': input_data['data'], 'data': input_data['data'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -221,14 +228,14 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': input_data['data'], 'data': input_data['data'],
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata **metadata
}, },
retry_policy=ANY, start_to_close_timeout=ANY) retry_policy=ANY, start_to_close_timeout=ANY)
@@ -261,7 +268,9 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'], 'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'}
} }
@@ -286,7 +295,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, { call(Activities.get_last_timestamp, {
'data': input_data['data'], 'data': input_data['data'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -294,14 +303,14 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': input_data['data'], 'data': input_data['data'],
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata **metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -310,7 +319,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
@@ -318,7 +327,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': 'transformed_data', 'data': 'transformed_data',
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -339,7 +348,9 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'mlflow_transform_filters': {'test': 'filter'}, 'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'}, 'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name', 'model_name': 'test_model_name',
'model_retention': '30', 'model_config': {
'retention': '30'
},
'path_priority': ['continue', 'repeat', 'stop'], 'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'}
} }
@@ -365,7 +376,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, { call(Activities.get_last_timestamp, {
'data': input_data['data'], 'data': input_data['data'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -373,14 +384,14 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'filters': input_data['input_filters'], 'filters': input_data['input_filters'],
'data': input_data['data'], 'data': input_data['data'],
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, },
retry_policy=ANY, start_to_close_timeout=ANY)]) retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_transform, { call(Activities.request_transform, {
'data': input_data['data'], 'data': input_data['data'],
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata **metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -389,7 +400,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, { call(Activities.mlflow_content_gate, {
@@ -397,13 +408,13 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': 'transformed_data', 'data': 'transformed_data',
'type': 'transform', 'type': 'transform',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.request_predict, { call(Activities.request_predict, {
'data': 'transformed_data', 'data': 'transformed_data',
'model_name': input_data['model_name'], 'model_name': input_data['model_name'],
'model_retention': input_data['model_retention'], 'model_config': input_data['model_config'],
**metadata **metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([ workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -412,7 +423,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict', 'type': 'predict',
'path_priority': input_data['path_priority'], 'path_priority': input_data['path_priority'],
**metadata **metadata,
}, retry_policy=ANY, start_to_close_timeout=ANY)]) }, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called()
@@ -429,7 +440,9 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
model = 'test_model' model = 'test_model'
last_timestamp = '2024-01-01' last_timestamp = '2024-01-01'
model_name = 'test_model_name' model_name = 'test_model_name'
model_retention = '30' model_config = {
'retention': '30'
}
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
@@ -440,7 +453,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
'model_id': model, 'model_id': model,
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_config': model_config
}, confidence, last_timestamp, "" }, confidence, last_timestamp, ""
) )
@@ -462,7 +475,9 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
model = 'test_model' model = 'test_model'
last_timestamp = '2024-01-01' last_timestamp = '2024-01-01'
model_name = 'test_model_name' model_name = 'test_model_name'
model_retention = '30' model_config = {
'retention': '30'
}
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
@@ -473,7 +488,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
'model_id': model, 'model_id': model,
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention 'model_config': model_config
}, confidence, last_timestamp, "" }, confidence, last_timestamp, ""
) )
@@ -506,7 +521,10 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
model = 'test_model' model = 'test_model'
last_timestamp = '2024-01-01' last_timestamp = '2024-01-01'
model_name = 'test_model_name' model_name = 'test_model_name'
model_retention = '30' model_config = {
'retention': '30'
}
prediction_store_policy = 'erl:1'
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
@@ -517,8 +535,9 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'model_id': model, 'model_id': model,
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention, 'model_config': model_config,
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, 'Prediction Process' }, confidence, last_timestamp, 'Prediction Process'
) )
@@ -535,11 +554,12 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'timestamp': last_timestamp, 'timestamp': last_timestamp,
'model_id': model, 'model_id': model,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention, 'model_config': model_config,
'schema': schema, 'schema': schema,
'table_name': table_name, 'table_name': table_name,
'comment': 'Prediction Process', 'comment': 'Prediction Process',
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
} }
) )
@@ -556,8 +576,10 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
model = 'test_model' model = 'test_model'
last_timestamp = '2024-01-01' last_timestamp = '2024-01-01'
model_name = 'test_model_name' model_name = 'test_model_name'
model_retention = '30' model_config = {
'retention': '30'
}
prediction_store_policy = 'erl:1'
# Act # Act
result = await prediction_process.path_flag_handler( result = await prediction_process.path_flag_handler(
data, path_flag, { data, path_flag, {
@@ -567,8 +589,9 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
'model_id': model, 'model_id': model,
'last_timestamp': last_timestamp, 'last_timestamp': last_timestamp,
'model_name': model_name, 'model_name': model_name,
'model_retention': model_retention, 'model_config': model_config,
'opc_output_config': {'test': 'config'} 'opc_output_config': {'test': 'config'},
'prediction_store_policy': prediction_store_policy
}, confidence, last_timestamp, "" }, confidence, last_timestamp, ""
) )

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

@@ -33,7 +33,11 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'schema': 'test_schema', 'schema': 'test_schema',
'table_name': 'test_table', 'table_name': 'test_table',
'opc_output_config': 'test_opc_output_config', 'opc_output_config': 'test_opc_output_config',
'datetime_columns': ['timestamp', 'created_at'] 'datetime_columns': ['timestamp', 'created_at'],
'prediction_store_policy': 'erl:1',
'model_config': {
'retention': '30'
}
} }
await predictions_batch.run(input_data) await predictions_batch.run(input_data)
@@ -72,9 +76,10 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
'POLICY': 'STOP' 'POLICY': 'STOP'
} }
}), }),
'model_retention': input_data.get('model_retention', 60), 'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}) 'opc_output_config': input_data.get('opc_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
} }
workflow_mock.execute_child_workflow.assert_has_calls([ workflow_mock.execute_child_workflow.assert_has_calls([