SIENTIAPDE-994

Implement new get_last_timestamp method in Gates class, refactor MLFlow activity methods to return only transformed data, and update PredictionsBatch and PredictionProcess workflows to utilize Activities module. Add detailed docstrings for new methods and enhance test coverage for get_last_timestamp functionality.
This commit is contained in:
vitor-aignosi
2025-05-12 11:36:21 -03:00
parent d09fb6ac5e
commit eb6d2dd79c
8 changed files with 540 additions and 83 deletions

View File

@@ -585,3 +585,19 @@ async def test_format_default_prediction(
result = await gates.format_default_prediction(input_data)
assert result == expected_output.to_dict()
@mark.asyncio
async def test_get_last_timestamp(
gates
):
input_data = {
'data': {
'variable': ['variable1', 'variable2'],
'value': [1, 2],
'timestamp': ['2021-01-01', '2021-01-02']
}
}
result = await gates.get_last_timestamp(input_data)
assert result == '2021-01-02'

View File

@@ -0,0 +1,422 @@
from unittest.mock import AsyncMock, patch, call
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
@fixture
def prediction_process():
return PredictionProcess()
@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 = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # 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)
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # 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([
call(Activities.input_gate, {
'filters': input_data['filters'],
'data': input_data['data']
})])
workflow_mock.execute_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([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
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_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': 'continue',
'data': 'predicted_data',
'prediction_confidence': 0.95,
'timestamp': '2024-01-01',
'model_id': 'test_model',
'model_name': 'test_model_name',
'model_retention': '30'
}
)
@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
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('stop', 0.95), # 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']}),
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})
])
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
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('repeat', 0.95), # input_gate
{'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data
('continue', 0.95), # 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([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_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([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
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
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # 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)
]
# 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([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_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([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
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
input_data = {
'data': {'test': 'data'},
'schema': 'test_schema',
'table_name': 'test_table',
'model': 'test_model',
'filters': {'test': 'filter'},
'model_name': 'test_model_name',
'model_retention': '30'
}
# Mock the activity responses
workflow_mock.execute_activity_method.side_effect = [
'2024-01-01', # get_last_timestamp
('continue', 0.95), # 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)
{'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict
('continue', 0.95), # 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([
call(Activities.input_gate, {
'filters': input_data['filters'], 'data': input_data['data']})])
workflow_mock.execute_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([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.mlflow_content_gate, {
'filters': input_data['filters'],
'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
'type': 'transform'
})])
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_retention': input_data['model_retention']
})])
workflow_mock.execute_activity_method.assert_has_calls([
call(Activities.mlflow_response_gate, {
'filters': input_data['filters'],
'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
'type': 'predict'
})])
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_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
)
# Assert
assert result is True
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_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_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.repeat_last_prediction,
{
'schema': schema,
'table_name': table_name,
'model': model
}
)
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_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
)
# Assert
assert result is True
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_called_once_with(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model,
'model_name': model_name,
'model_retention': model_retention
}
)
@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_retention = '30'
# Act
result = await prediction_process.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
)
# Assert
assert result is False
workflow_mock.execute_activity_method.assert_not_called()
workflow_mock.execute_child_workflow.assert_not_called()