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

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

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
-
-        return response_data, timestamp
 
     @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'])

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.activities import Activities
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
     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(
+        input_data['data'] = data
-            Gates.input_gate,
-            {
-                'filters': input_data['filters'],
-                'data': data
-            }
-        )
 
-        if path_flag == 'stop':
+        await workflow.execute_child_workflow(
-            return
+            'prediction_process', input_data)
-
-        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

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']

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.activities import Activities
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
     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(
+        response_data = await workflow.execute_activity_method(
-            MLFlow.transform_data,
+            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,

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'

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