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()