From eb6d2dd79c573c43be25ad66357753f5c1c4799d Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 12 May 2025 11:36:21 -0300 Subject: [PATCH] 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. --- laborious/activities/gates.py | 13 + laborious/activities/mlflow.py | 12 +- laborious/workflows/predictions_batch.py | 67 +-- .../format_and_export_prediction.py | 23 + .../sub_workflows/prediction_process.py | 70 ++- tests/laborious/activities/test_gates.py | 16 + .../test_format_and_export_prediction.py | 0 .../subworkflows/test_prediction_process.py | 422 ++++++++++++++++++ 8 files changed, 540 insertions(+), 83 deletions(-) rename tests/laborious/workflows/{subworkflows.py => subworkflows}/test_format_and_export_prediction.py (100%) create mode 100644 tests/laborious/workflows/subworkflows/test_prediction_process.py diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index a2fad98..4d9eb26 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -207,3 +207,16 @@ class Gates(BaseActivity): 'prediction_status': ['Bad'], 'comment': [input_data['comment']] }).to_dict() + + @activity.defn(name="get_last_timestamp") + async def get_last_timestamp(self, input_data: dict[str, Any]) -> str: + """ + Gets the last timestamp of the data. + Args: + input_data (dict): The input data. Contains: + data (dict[str, Any]): The data to get the last timestamp from. + Returns: + str: The last timestamp of the data. + """ + data = DataFrame(input_data['data']) + return max(data['timestamp'].values.tolist()) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 7c785b5..d1b8d4c 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -26,7 +26,7 @@ class MLFlow(BaseActivity): ) @activity.defn(name="request_transform") - async def request_transform(self, input_data: dict[str, Any]) -> tuple[dict[str, Any], str]: + async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Access MLFlow model to get the transformed data. Args: @@ -35,7 +35,7 @@ class MLFlow(BaseActivity): model_name (str): The name of the model. model_retention (int): The retention of the model. Returns: - tuple[dict[str, Any], str]: The transformed data and the latest timestamp of the data. + dict[str, Any]: The transformed data. """ self.logger.info('Transforming data...') data = DataFrame(input_data['data']) @@ -54,12 +54,10 @@ class MLFlow(BaseActivity): response_data = self.model_monitoring_repository.transform( model_name, data, model_retention) - timestamp = max(data['timestamp'].values.tolist()) - - return response_data, timestamp + return response_data @activity.defn(name="request_predict") - async def request_predict(self, input_data: dict[str, Any]) -> tuple[dict[str, Any], str]: + async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Access MLFlow model to get the predicted data. Args: @@ -68,7 +66,7 @@ class MLFlow(BaseActivity): model_name (str): The name of the model. model_retention (int): The retention of the model. Returns: - tuple[dict[str, Any], str]: The predicted data and the latest timestamp of the data. + dict[str, Any]: The predicted data. """ self.logger.info('Predicting data...') data = DataFrame(input_data['data']) diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index c3c54b0..6c137da 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -1,10 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.postgres import Postgres - from laborious.activities.mlflow import MLFlow - from laborious.activities.gates import Gates - from laborious.activities.opc import OPC + from laborious.activities.activities import Activities from typing import Any @@ -14,7 +11,7 @@ class PredictionsBatch(): async def run(self, input_data: dict[str, Any]): await workflow.execute_activity_method( - Postgres.prepare_activity, + Activities.prepare_activity, { 'schedule_name': input_data['schedule_name'], 'model_name': input_data['model_name'], @@ -23,63 +20,11 @@ class PredictionsBatch(): ) data = await workflow.execute_activity_method( - Postgres.load_custom_query, + Activities.load_custom_query, input_data['query'] ) - path_flag, confidence = await workflow.execute_activity_method( - Gates.input_gate, - { - 'filters': input_data['filters'], - 'data': data - } - ) + input_data['data'] = data - if path_flag == 'stop': - return - - if path_flag == 'continue': - # repeat last prediction - await workflow.execute_activity_method( - Postgres.repeat_last_prediction, - { - 'schema': input_data['schema'], - 'table_name': input_data['table_name'], - 'model': input_data['model'] - } - ) - return - - response_data, last_timestamp = await workflow.execute_activity_method( - MLFlow.transform_data, - { - 'data': data, - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - } - ) - - path_flag, confidence = await workflow.execute_activity_method( - Gates.mlflow_gate, - { - 'filters': input_data['filters'], - 'data': response_data, - 'type': 'transform' - } - ) - - if path_flag == 'stop': - return - - if path_flag is None: - # procced with prediction - response_data = await workflow.execute_activity_method( - MLFlow.request_predict, - { - 'data': response_data, - 'model_name': input_data['model_name'], - 'model_retention': input_data['model_retention'] - } - ) - - path_flag + await workflow.execute_child_workflow( + 'prediction_process', input_data) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 605b69d..553c9f8 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -9,6 +9,29 @@ with workflow.unsafe.imports_passed_through(): class FormatAndExportPrediction(): @workflow.run async def run(self, input_data: dict[str, Any]): + """ + This workflow formats and exports predictions based on path_flag: + - If path_flag is None: formats prediction using input data, timestamp, model_id and confidence + - If path_flag exists: creates default prediction with timestamp, model_id, confidence and comment + Finally exports formatted prediction to postgres table + Args: + input_data(dict[str, Any]): The input data for the workflow. Contains the following keys: + - path_flag(str): The path flag to determine the type of prediction to format + - data(dict[str, Any]): The data to format + - prediction_confidence(float): The prediction confidence to be registered + - timestamp(str): The timestamp of the prediction, synchronized with the data + - model_id(str): The model id of the prediction + - model_name(str): The model name of the prediction + - model_retention(str): The model retention of the prediction + - comment(str): The comment to be registered + - schema(str): The schema of the prediction + - table_name(str): The table name of the prediction + - opc_servers(list[str]): The opc servers of the prediction + - opc_output_config(dict[str, Any]): The opc output config of the prediction + + Returns: + bool: True if the workflow was successful, False otherwise. + """ path_flag = input_data['path_flag'] data = input_data['data'] prediction_confidence = input_data['prediction_confidence'] diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 60871fe..05b3350 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -1,10 +1,7 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): - from laborious.activities.postgres import Postgres - from laborious.activities.mlflow import MLFlow - from laborious.activities.gates import Gates - from laborious.activities.opc import OPC + from laborious.activities.activities import Activities from typing import Any @@ -20,8 +17,15 @@ class PredictionProcess(): model_name = input_data['model_name'] model_retention = input_data['model_retention'] + last_timestamp = await workflow.execute_activity_method( + Activities.get_last_timestamp, + { + 'data': data + } + ) + path_flag, confidence = await workflow.execute_activity_method( - Gates.input_gate, + Activities.input_gate, { 'filters': input_data['filters'], 'data': data @@ -29,12 +33,13 @@ class PredictionProcess(): ) if await self.path_flag_handler( - data, path_flag, confidence, schema, table_name, model + data, path_flag, confidence, schema, table_name, + model, last_timestamp, model_name, model_retention ): return - response_data, last_timestamp = await workflow.execute_activity_method( - MLFlow.transform_data, + response_data = await workflow.execute_activity_method( + Activities.request_transform, { 'data': data, 'model_name': model_name, @@ -43,7 +48,7 @@ class PredictionProcess(): ) path_flag, confidence = await workflow.execute_activity_method( - Gates.mlflow_gate, + Activities.mlflow_response_gate, { 'filters': filters, 'data': response_data, @@ -52,12 +57,28 @@ class PredictionProcess(): ) if await self.path_flag_handler( - data, path_flag, confidence, schema, table_name, model + data, path_flag, confidence, schema, table_name, + model, last_timestamp, model_name, model_retention + ): + return + + path_flag, confidence = await workflow.execute_activity_method( + Activities.mlflow_content_gate, + { + 'filters': filters, + 'data': response_data, + 'type': 'transform' + } + ) + + if await self.path_flag_handler( + data, path_flag, confidence, schema, table_name, + model, last_timestamp, model_name, model_retention ): return response_data = await workflow.execute_activity_method( - MLFlow.request_predict, + Activities.request_predict, { 'data': response_data, 'model_name': model_name, @@ -66,7 +87,7 @@ class PredictionProcess(): ) path_flag, confidence = await workflow.execute_activity_method( - Gates.mlflow_gate, + Activities.mlflow_response_gate, { 'filters': filters, 'data': response_data, @@ -75,7 +96,8 @@ class PredictionProcess(): ) if await self.path_flag_handler( - data, path_flag, confidence, schema, table_name, model + data, path_flag, confidence, schema, table_name, + model, last_timestamp, model_name, model_retention ): return @@ -94,14 +116,32 @@ class PredictionProcess(): async def path_flag_handler(self, data: dict[str, Any], path_flag: str, confidence: int, schema: str, table_name: str, - model: str): + model: str, last_timestamp: str, model_name: str, + model_retention: str): + """ + This function handles the path flag and the confidence of the prediction. + It returns True if the prediction should be stopped. If path_flag is 'repeat', it repeats the last prediction. + If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop', it stops the prediction process. + Args: + data (dict[str, Any]): The data to be used for the prediction. + path_flag (str): The path flag to determine the type of prediction to format + confidence (int): The confidence of the prediction + schema (str): The schema of the prediction + table_name (str): The table name of the prediction + model (str): The model id of the prediction + last_timestamp (str): The timestamp of the last prediction + model_name (str): The model name of the prediction + model_retention (str): The model retention of the prediction + Returns: + bool: True if the prediction should be stopped, False otherwise. + """ if path_flag == 'stop': return True elif path_flag == 'repeat': # repeat last prediction await workflow.execute_activity_method( - Postgres.repeat_last_prediction, + Activities.repeat_last_prediction, { 'schema': schema, 'table_name': table_name, diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index ff84955..87c8254 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -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' diff --git a/tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py similarity index 100% rename from tests/laborious/workflows/subworkflows.py/test_format_and_export_prediction.py rename to tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py new file mode 100644 index 0000000..67c81e1 --- /dev/null +++ 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()