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

SIENTIAPDE-1110

Browse Source 
Refactor notification handling in Gates and OPC activities to use send_notification method with metadata integration, enhancing error reporting and traceability.
This commit is contained in:
Vitor Pimentel dos Santos
2025-07-04 15:19:14 -03:00
parent 58216a91ce
commit 50e6479c49
12 changed files with 293 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -10,11 +10,22 @@ 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",
@@ -35,7 +46,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'data': input_data['data'],
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
'prediction_confidence': input_data['prediction_confidence'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -46,7 +58,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -59,7 +72,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value
'data': workflow_mock.execute_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -74,6 +88,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
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",
@@ -95,7 +110,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence'],
'comment': input_data['comment']
'comment': input_data['comment'],
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -107,7 +123,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
Activities.write_opc_data,
{
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_local_activity_method.return_value
'data': workflow_mock.execute_local_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY
@@ -120,7 +137,8 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': workflow_mock.execute_activity_method.return_value
'data': workflow_mock.execute_activity_method.return_value,
**metadata
},
retry_policy=ANY,
start_to_close_timeout=ANY

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

@@ -9,12 +9,23 @@ def prediction_process():
return PredictionProcess()
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.prediction_process.workflow", new_callable=AsyncMock)
async def test_run(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=False)
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -49,51 +60,61 @@ async def test_run(workflow_mock, prediction_process):
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
retry_policy=ANY, start_to_close_timeout=ANY)])
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
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']
'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.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**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']
'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': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'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.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**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']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'metadata': metadata,
'path_flag': 'continue',
'data': 'predicted_data',
'prediction_confidence': 0.95,
@@ -115,6 +136,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
prediction_process.path_flag_handler = AsyncMock(return_value=True)
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -141,11 +163,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {
'data': input_data['data']}, retry_policy=ANY, start_to_close_timeout=ANY),
'data': input_data['data'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY),
call(Activities.input_gate, {
'filters': input_data['input_filters'],
'data': input_data['data'],
'path_priority': input_data['path_priority']}, retry_policy=ANY, start_to_close_timeout=ANY)
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -156,6 +182,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -183,19 +210,26 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 4
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
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']},
'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.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']},
'model_retention': input_data['model_retention'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
@@ -203,7 +237,8 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -216,6 +251,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
side_effect=[False, False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -247,33 +283,41 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
assert workflow_mock.execute_local_activity_method.call_count == 5
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
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']},
'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.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**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']
'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': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -285,6 +329,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
side_effect=[False, False, False, True])
# Arrange
input_data = {
'metadata': metadata,
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
@@ -317,46 +362,56 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
# Assert
assert workflow_mock.execute_local_activity_method.call_count == 7
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']},
call(Activities.get_last_timestamp, {
'data': input_data['data'],
**metadata
},
retry_policy=ANY, start_to_close_timeout=ANY)])
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']},
'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.request_transform, {
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**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']
'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': 'transformed_data',
'type': 'transform',
'path_priority': input_data['path_priority']
'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.request_predict, {
'data': 'transformed_data',
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
'model_retention': input_data['model_retention'],
**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']
'path_priority': input_data['path_priority'],
**metadata
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -378,6 +433,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -410,6 +466,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -424,6 +481,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.repeat_last_prediction,
{
**metadata,
'schema': schema,
'table_name': table_name,
'model': model,
@@ -452,6 +510,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
'metadata': metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,
@@ -468,6 +527,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'metadata': metadata,
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
@@ -500,6 +560,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, {
**metadata,
'schema': schema,
'table_name': table_name,
'model_id': model,

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

@@ -9,6 +9,16 @@ def predictions_batch() -> PredictionsBatch:
return PredictionsBatch()
metadata = {
"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):
@@ -29,27 +39,17 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
Activities.load_custom_query,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
**metadata,
'query': input_data['query'],
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query'],
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],