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.
2025-05-12 11:36:21 -03:00
parent d09fb6ac5e
commit eb6d2dd79c
8 changed files with 540 additions and 83 deletions
--- a/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py
+++ b/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py
--- a/tests/laborious/workflows/subworkflows/test_prediction_process.py
+++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py
@@ -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()