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

SIENTIAPDE-994

Browse Source 
Refactor and enhance the laborious workflow and utilities

- Removed outdated test file `test_predictions_batch.py` from workflows.
- Added `input_sample.json` for standardized input configuration.
- Introduced `connectors_config.py` to manage database and service configurations.
- Implemented a logging utility in `logger.py` for consistent logging across the application.
- Created `policies.py` to define retry policies for workflows.
- Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`.
- Added extensive tests for `OpcRepository` in `test_opc_repository.py`.
- Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
This commit is contained in:
vitor-aignosi
2025-05-23 17:34:47 -03:00
parent 5fe552410b
commit 67fe4afaa6
30 changed files with 1385 additions and 765 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,4 +1,4 @@
from unittest.mock import call, patch, AsyncMock
from unittest.mock import call, patch, AsyncMock, ANY
from pytest import mark, fixture
from laborious.activities.activities import Activities
@@ -28,7 +28,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_prediction,
{
@@ -36,7 +36,9 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': input_data['prediction_confidence']
}
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
@@ -44,8 +46,10 @@ 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_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)])
workflow_mock.execute_activity_method.assert_has_calls([
call(
@@ -53,12 +57,15 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
{
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value
}
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1
@mark.asyncio
@@ -80,7 +87,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
await format_and_export_prediction.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.format_default_prediction,
{
@@ -88,7 +95,9 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
'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([
@@ -97,8 +106,10 @@ 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_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
workflow_mock.execute_activity_method.assert_has_calls([
@@ -107,9 +118,12 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
{
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'],
'data': workflow_mock.execute_activity_method.return_value
}
'data': workflow_mock.execute_local_activity_method.return_value
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
])
assert workflow_mock.execute_activity_method.call_count == 3
assert workflow_mock.execute_activity_method.call_count == 2
assert workflow_mock.execute_local_activity_method.call_count == 1

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

@@ -1,4 +1,4 @@
from unittest.mock import AsyncMock, patch, call
from unittest.mock import AsyncMock, patch, call, ANY
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@@ -18,65 +18,78 @@ async def test_run(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'},
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
# mlflow_response_gate (transform)
('continue', 0.95, "Error"),
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict)
# mlflow_response_gate (predict)
('continue', 0.95, "Error"),
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
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)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'],
'data': input_data['data']
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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_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']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, 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['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, 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['filters'],
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
'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(
'format_and_export_prediction',
@@ -85,9 +98,10 @@ async def test_run(workflow_mock, prediction_process):
'data': 'predicted_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 'test_model',
'model_id': 1,
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'opc_output_config': input_data['opc_output_config']
}
)
@@ -101,27 +115,34 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95), # input_gate
('stop', 0.95, "Input data with bad quality"), # input_gate
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 2
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']}),
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),
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})
'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_child_workflow.assert_not_called()
@@ -135,42 +156,53 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95), # input_gate
('repeat', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 4
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
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']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'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_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']
})])
workflow_mock.execute_activity_method.assert_has_calls([
'model_retention': input_data['model_retention']},
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['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)
])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -184,49 +216,61 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
# 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_activity_method.call_count == 5
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
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']},
retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'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_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']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, 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['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
'type': 'transform',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -240,63 +284,75 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_id': 1,
'input_filters': {'test': 'filter'},
'mlflow_transform_filters': {'test': 'filter'},
'mlflow_predict_filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
'model_retention': '30',
'path_priority': ['continue', 'repeat', 'stop'],
'opc_output_config': {'test': 'config'}
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
workflow_mock.execute_local_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # input_gate
('continue', 0.95, "Input data with bad quality"), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # mlflow_response_gate (transform)
('continue', 0.95), # mlflow_content_gate (transform)
('continue', 0.95, "Error"), # mlflow_response_gate (transform)
# mlflow_content_gate (transform)
('continue', 0.95, "Transformed data not passed the content filter"),
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # mlflow_response_gate (predict)
('continue', 0.95, "Error"), # mlflow_response_gate (predict)
]
# Act
await prediction_process.run(input_data)
# Assert
assert workflow_mock.execute_activity_method.call_count == 7
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.get_last_timestamp, {'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
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)])
workflow_mock.execute_local_activity_method.assert_has_calls([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_activity_method.assert_has_calls([
'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_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']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, 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['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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, {
'filters': input_data['filters'],
'filters': input_data['mlflow_transform_filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
'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.request_predict, {
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'model_name': input_data['model_name'],
'model_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
}, 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['filters'],
'filters': input_data['mlflow_predict_filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
'type': 'predict',
'path_priority': input_data['path_priority']
}, retry_policy=ANY, start_to_close_timeout=ANY)])
workflow_mock.execute_child_workflow.assert_not_called()
@@ -305,7 +361,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'stop'
path_flag = 'STOP'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
@@ -317,12 +373,12 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_local_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()
@@ -343,7 +399,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert
@@ -353,8 +409,10 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
{
'schema': schema,
'table_name': table_name,
'model': model
}
'model_id': model
},
retry_policy=ANY,
start_to_close_timeout=ANY
)
workflow_mock.execute_child_workflow.assert_not_called()
@@ -364,7 +422,7 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Arrange
data = {'test': 'data'}
path_flag = 'continue'
path_flag = 'CONTINUE'
confidence = 0.95
schema = 'test_schema'
table_name = 'test_table'
@@ -376,7 +434,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, 'Prediction Process'
)
# Assert
@@ -391,7 +449,10 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': 'Prediction Process'
}
)
@@ -413,7 +474,7 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
model, last_timestamp, model_name, model_retention, ""
)
# Assert

48
tests/laborious/workflows/subworkflows/test_predictions_batch.py
View File

@@ -1,48 +0,0 @@
from unittest.mock import AsyncMock, 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()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_activity_method.return_value = {
'data': 'test_data'
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'query': 'SELECT * FROM test'
}
await predictions_batch.run(input_data)
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id']
}
)
])
workflow_mock.execute_activity_method.assert_has_calls([
call(
Activities.load_custom_query,
input_data['query']
)
])
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', input_data)
])

68
tests/laborious/workflows/test_predictions_batch.py Normal file
View File

@@ -0,0 +1,68 @@
from unittest.mock import AsyncMock, call, patch, ANY
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()
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {
'data': 'test_data'
}
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',
'opc_output_config': 'test_opc_output_config'
}
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls([
call(
Activities.prepare_activity,
{
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch'
},
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 = {
'data': {'data': 'test_data'},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get('input_filters', {}),
'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {}),
'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {}),
'model_retention': input_data.get('model_retention', 60),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
}
workflow_mock.execute_child_workflow.assert_has_calls([
call(
'prediction_process', prediction_input)
])