SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows

This commit includes several changes:

- Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability.
- Improved error handling and logging in gates.py and mlflow.py.
- Added input validation and filtering in gates.py to ensure data quality.
- Enhanced prediction formatting and storage policy management in gates.py.
- Updated metrics.py to use consistent naming conventions and labels.
- Refactored connectors_config.py to use type hints and improve code clarity.
- Updated conditional and MLFlow filters for better data quality checks.
- Improved model repository logic for retraining and updating models.
- Enhanced worker.py to include SDK metrics and improved error handling.
- Refactored workflows for better modularity and error handling.
- Updated tests to reflect the changes and improve test coverage.

model_manager/activities/activities.py

@@ -1,12 +1,14 @@
-from temporalio import activity, workflow
+from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from sientia_do.temporal.activities.postgres import Postgres
+    from typing import Any
+
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import Logger
-    from model_manager.activities.mlflow import MLFlow
+    from sientia_do.temporal.activities.postgres import Postgres
+
     from model_manager.activities.gates import Gates
-    from typing import Any
+    from model_manager.activities.mlflow import MLFlow
 
 
 class Activities(Postgres, MLFlow, Gates):
@@ -29,11 +31,13 @@ class Activities(Postgres, MLFlow, Gates):
         notification_handler (NotificationHandler): Notification management instance
     """
 
-    def __init__(self,
-                 postgres_config: dict[str, Any],
-                 mlflow_config: dict[str, Any],
-                 logger: Logger,
-                 notification_handler: NotificationHandler):
+    def __init__(
+        self,
+        postgres_config: dict[str, Any],
+        mlflow_config: dict[str, Any],
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
         """
         Initialize the Activities orchestrator with all required configurations.
 
@@ -52,25 +56,30 @@ class Activities(Postgres, MLFlow, Gates):
             Exception: If any parent class initialization fails
         """
         # Initialize parent classes
-        Postgres.__init__(self, host=postgres_config['host'],
-                          port=postgres_config['port'],
-                          user=postgres_config['user'],
-                          password=postgres_config['password'],
-                          dbname=postgres_config['dbname'],
-                          min_connections=postgres_config['min_connections'],
-                          max_connections=postgres_config['max_connections'],
-                          logger=logger,
-                          notification_handler=notification_handler)
+        Postgres.__init__(
+            self,
+            host=postgres_config['host'],
+            port=postgres_config['port'],
+            user=postgres_config['user'],
+            password=postgres_config['password'],
+            dbname=postgres_config['dbname'],
+            min_connections=postgres_config['min_connections'],
+            max_connections=postgres_config['max_connections'],
+            logger=logger,
+            notification_handler=notification_handler,
+        )
 
-        MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
-                        mlflow_port=mlflow_config['port'],
-                        mlflow_username=mlflow_config['username'],
-                        mlflow_password=mlflow_config['password'],
-                        logger=logger,
-                        notification_handler=notification_handler)
+        MLFlow.__init__(
+            self,
+            mlflow_host=mlflow_config['host'],
+            mlflow_port=mlflow_config['port'],
+            mlflow_username=mlflow_config['username'],
+            mlflow_password=mlflow_config['password'],
+            logger=logger,
+            notification_handler=notification_handler,
+        )
 
-        Gates.__init__(self, logger=logger,
-                       notification_handler=notification_handler)
+        Gates.__init__(self, logger=logger, notification_handler=notification_handler)
 
     async def shutdown(self):
         """

model_manager/activities/gates.py

@@ -1,53 +1,42 @@
 from temporalio import activity, workflow
 
-
 with workflow.unsafe.imports_passed_through():
     import traceback
+    from typing import Any
+
+    from pandas import DataFrame
+    from sientia_do.formatters import create_sample_dict
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.observability.logger import Logger
+    from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
-    from sientia_do.formatters import create_sample_dict
-    from model_manager.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
 
+    from model_manager import metrics
     from model_manager.utils.filters.conditional_filters import (
         filter_empty_data,
-        filter_specific_variables_null_values
+        filter_specific_variables_null_values,
     )
-    from pandas import DataFrame
-    from model_manager import metrics
+    from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
 
 # Input filter function mappings
 input_filter_functions = {
     'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
     'EMPTY_DATA': filter_empty_data,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 2,
-        'REPEAT': -1
-    }
+    'path_confidence': {'STOP': -1, 'CONTINUE': 2, 'REPEAT': -1},
 }
 
 # MLFlow response filter function mappings
 mlflow_response_filter_functions = {
     'API_ERROR': api_error_filter,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 10,
-        'REPEAT': -1
-    },
+    'path_confidence': {'STOP': -1, 'CONTINUE': 10, 'REPEAT': -1},
 }
 
 # MLFlow content filter function mappings
 mlflow_content_filter_functions = {
     'NAN_VALUES': nan_values_filter,
     'EMPTY_DATA': filter_empty_data,
-    'path_confidence': {
-        'STOP': -1,
-        'CONTINUE': 18,
-        'REPEAT': -1
-    }
+    'path_confidence': {'STOP': -1, 'CONTINUE': 18, 'REPEAT': -1},
 }
 
 
@@ -82,10 +71,9 @@ class Gates(BaseActivity):
         Raises:
             Exception: If BaseActivity initialization fails
         """
-        BaseActivity.__init__(
-            self, logger, notification_handler, set_error_counter=True)
+        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
 
-    @activity.defn(name="input_gate")
+    @activity.defn(name='input_gate')
     async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Apply input data quality filters and validation.
@@ -121,7 +109,7 @@ class Gates(BaseActivity):
         """
         metadata = input_data['metadata']
 
-        self.info("Performing input gate...", metadata)
+        self.info('Performing input gate...', metadata)
 
         filters = input_data['filters']
         data = DataFrame(input_data['data'])
@@ -129,40 +117,42 @@ class Gates(BaseActivity):
 
         filter_output = []
 
-        self.debug(f"Input data: {data.head(5).to_string()}", metadata)
-        self.debug(f"Filters: {filters}", metadata)
+        self.debug(f'Input data: {data.head(5).to_string()}', metadata)
+        self.debug(f'Filters: {filters}', metadata)
 
         # Apply each configured filter
         for fil, config in filters.items():
             if fil not in input_filter_functions:
-                self.error(f"Filter {fil} not found", metadata)
+                self.error(f'Filter {fil} not found', metadata)
                 continue
             try:
                 if input_filter_functions[fil](data, config['config']):
-                    self.debug(
-                        f"Data not passed the input filter {fil}:{config}", metadata)
+                    self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
                     filter_output.append(config['policy'])
-            except Exception as e:
+            except Exception as e:  # noqa: BLE001
                 trace = traceback.format_exc()
                 self.send_notification(
                     metadata=metadata,
-                    notification_id=f"INTPUT_GATE_ERROR__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="input_gate",
+                    notification_id=f'INTPUT_GATE_ERROR__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='input_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(f"Input gate result: {path_flag}", metadata)
-                return path_flag, input_filter_functions['path_confidence'][path_flag], \
-                    "Input data with bad quality"
+                self.info(f'Input gate result: {path_flag}', metadata)
+                return (
+                    path_flag,
+                    input_filter_functions['path_confidence'][path_flag],
+                    'Input data with bad quality',
+                )
 
-        self.info("Nothing was filtered by the input gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the input gate', metadata)
+        return None, 0, ''
 
-    @activity.defn(name="mlflow_response_gate")
+    @activity.defn(name='mlflow_response_gate')
     async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Validate MLFlow API response quality and integrity.
@@ -197,7 +187,7 @@ class Gates(BaseActivity):
             Exception: If response validation fails or configuration is invalid
         """
         metadata = input_data['metadata']
-        self.info("Performing mlflow response gate...", metadata)
+        self.info('Performing mlflow response gate...', metadata)
 
         filters = input_data['filters']
         data = input_data['data']
@@ -206,14 +196,13 @@ class Gates(BaseActivity):
 
         filter_output = []
 
-        self.debug(
-            f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata)
-        self.debug(f"Filters: {filters}", metadata)
+        self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}', metadata)
+        self.debug(f'Filters: {filters}', metadata)
 
         comments = []
         for fil, config in filters.items():
             if fil not in mlflow_response_filter_functions:
-                self.error(f"Filter {fil} not found", metadata)
+                self.error(f'Filter {fil} not found', metadata)
                 continue
             try:
                 if mlflow_response_filter_functions[fil](data, config):
@@ -221,34 +210,36 @@ class Gates(BaseActivity):
                     comments.append(data['content']['message'])
                     self.send_notification(
                         metadata=metadata,
-                        notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
+                        notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
                         message=data['content']['message'],
-                        block="mlflow_gate",
+                        block='mlflow_gate',
                         level=NotificationLevel.ERROR,
-                        attachment_content=data['content']['traceback']
+                        attachment_content=data['content']['traceback'],
                     )
-            except Exception as e:
+            except Exception as e:  # noqa: BLE001
                 trace = traceback.format_exc()
                 self.send_notification(
                     metadata=metadata,
-                    notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="mlflow_gate",
+                    notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='mlflow_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(
-                    f"Mlflow response gate result: {path_flag}", metadata)
-                return path_flag, mlflow_response_filter_functions['path_confidence'][path_flag], \
-                    ", ".join(comments)
+                self.info(f'Mlflow response gate result: {path_flag}', metadata)
+                return (
+                    path_flag,
+                    mlflow_response_filter_functions['path_confidence'][path_flag],
+                    ', '.join(comments),
+                )
 
-        self.info("Nothing was filtered by the mlflow response gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the mlflow response gate', metadata)
+        return None, 0, ''
 
-    @activity.defn(name="mlflow_content_gate")
+    @activity.defn(name='mlflow_content_gate')
     async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
         Validate MLFlow prediction content quality and integrity.
@@ -283,7 +274,7 @@ class Gates(BaseActivity):
             Exception: If content validation fails or configuration is invalid
         """
         metadata = input_data['metadata']
-        self.info("Performing mlflow content gate...", metadata)
+        self.info('Performing mlflow content gate...', metadata)
 
         filters = input_data['filters']
         data = DataFrame(input_data['data'])
@@ -292,8 +283,8 @@ class Gates(BaseActivity):
 
         filter_output = []
 
-        self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
-        self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata)
+        self.debug(f'Input data:\n {data.head(5).to_string()}', metadata)
+        self.debug(f'Filters: \n {create_sample_dict(filters)}', metadata)
 
         for fil, config in filters.items():
             if fil not in mlflow_content_filter_functions:
@@ -303,36 +294,38 @@ class Gates(BaseActivity):
                     filter_output.append(config['policy'])
                     self.send_notification(
                         metadata=metadata,
-                        notification_id=f"{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}",
-                        message=f"Data not passed the content filter {fil}:{config}",
-                        block="mlflow_gate",
+                        notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
+                        message=f'Data not passed the content filter {fil}:{config}',
+                        block='mlflow_gate',
                         level=NotificationLevel.WARNING,
-                        attachment_content=data.to_string()
+                        attachment_content=data.to_string(),
                     )
-            except Exception as e:
+            except Exception as e:  # noqa: BLE001
                 trace = traceback.format_exc()
                 self.send_notification(
                     metadata=metadata,
-                    notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
-                    message=f"Error in filter {fil}:{config}: \n {e}",
-                    block="mlflow_gate",
+                    notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
+                    message=f'Error in filter {fil}:{config}: \n {e}',
+                    block='mlflow_gate',
                     level=NotificationLevel.ERROR,
-                    attachment_content=trace
+                    attachment_content=trace,
                 )
 
         for path_flag in path_priority:
             if path_flag in filter_output:
-                self.info(
-                    f"Mlflow content gate result: {path_flag}", metadata)
-                return path_flag, mlflow_content_filter_functions['path_confidence'][path_flag], \
-                    "Transformed data not passed the content filter"
+                self.info(f'Mlflow content gate result: {path_flag}', metadata)
+                return (
+                    path_flag,
+                    mlflow_content_filter_functions['path_confidence'][path_flag],
+                    'Transformed data not passed the content filter',
+                )
 
-        self.info("Nothing was filtered by the mlflow content gate", metadata)
-        return None, 0, ""
+        self.info('Nothing was filtered by the mlflow content gate', metadata)
+        return None, 0, ''
 
-    def get_prediction_store_policy(self,
-                                    prediction_store_policy: str,
-                                    metadata: dict[str, Any]) -> tuple[str, int]:
+    def get_prediction_store_policy(
+        self, prediction_store_policy: str, metadata: dict[str, Any]
+    ) -> tuple[str, int]:
         """
         Parse and validate prediction store policy configuration.
 
@@ -358,7 +351,9 @@ class Gates(BaseActivity):
 
         if len(policy_elements) < 2:
             self.error(
-                f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
+                f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
+                metadata,
+            )
             return 'lts', 1
 
         policy_type = policy_elements[0]
@@ -366,14 +361,20 @@ class Gates(BaseActivity):
 
         # If the policy_type is not lts or erl, we use the default policy
         # If the policty_value is not a number or 0, we use the default policy
-        if policy_type not in ['lts', 'erl'] or not policy_value.isdigit() or int(policy_value) == 0:
+        if (
+            policy_type not in ['lts', 'erl']
+            or not policy_value.isdigit()
+            or int(policy_value) == 0
+        ):
             self.error(
-                f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
+                f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
+                metadata,
+            )
             return 'lts', 1
 
         return policy_type, int(policy_value)
 
-    @activity.defn(name="format_prediction")
+    @activity.defn(name='format_prediction')
     async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Format prediction data according to configured storage policies.
@@ -400,7 +401,7 @@ class Gates(BaseActivity):
         """
         metadata = input_data['metadata']
         prediction_store_policy = input_data['prediction_store_policy']
-        self.info("Formatting prediction...", metadata)
+        self.info('Formatting prediction...', metadata)
 
         data = DataFrame(input_data['data'])
 
@@ -408,48 +409,45 @@ class Gates(BaseActivity):
         data['timestamp'] = data.index
         data = data.reset_index(drop=True)
 
-        self.debug(
-            f"Prediction store policy: {prediction_store_policy}", metadata)
-        self.debug(f"Prediction data: {data.head(5).to_string()}", metadata)
+        self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
+        self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
 
         policy_type, policy_value = self.get_prediction_store_policy(
-            prediction_store_policy, metadata)
+            prediction_store_policy, metadata
+        )
 
         # If data has no timestamp, we use the default timestamp and not sort the data
         self.info(
-            f"Sorting data by timestamp and applying policy: {policy_type}:{policy_value}", metadata)
+            f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
+        )
 
         # If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
         if policy_type == 'lts':
-            self.debug(
-                "Sorting data by timestamp descending", metadata)
+            self.debug('Sorting data by timestamp descending', metadata)
             data = data.sort_values(by='timestamp', ascending=False)
         # If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
         elif policy_type == 'erl':
-            self.debug(
-                "Sorting data by timestamp ascending", metadata)
+            self.debug('Sorting data by timestamp ascending', metadata)
             data = data.sort_values(by='timestamp', ascending=True)
         else:
-            self.error(
-                f"Invalid policy type: {policy_type}, using default policy", metadata)
-            raise ValueError(
-                f"Invalid policy type: {policy_type}")
+            self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
+            raise ValueError(f'Invalid policy type: {policy_type}')
 
         data = data.head(int(policy_value))
 
         data['model_id'] = input_data['model_id']
         data['prediction_confidence'] = input_data['prediction_confidence']
         data['prediction_status'] = 'Good'
-        data['comments'] = ""
+        data['comments'] = ''
         data = data.sort_values(by='timestamp', ascending=False)
         data = data.reset_index(drop=True)
 
-        self.info(f"Prediction formatted: {len(data)} rows", metadata)
-        self.debug(f"Prediction data: {data.head(5).to_string()}", metadata)
+        self.info(f'Prediction formatted: {len(data)} rows', metadata)
+        self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
 
         return data.to_dict()
 
-    @activity.defn(name="format_default_prediction")
+    @activity.defn(name='format_default_prediction')
     async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Create and format default prediction data for error conditions.
@@ -477,22 +475,24 @@ class Gates(BaseActivity):
         """
 
         metadata = input_data['metadata']
-        self.debug("Formatting default prediction...", metadata)
+        self.debug('Formatting default prediction...', metadata)
 
-        data = DataFrame({
-            'prediction': [0],
-            'response_time': [0],
-            'timestamp': [input_data['timestamp']],
-            'model_id': [input_data['model_id']],
-            'prediction_confidence': [input_data['prediction_confidence']],
-            'prediction_status': ['Bad'],
-            'comments': [input_data['comment']]
-        })
+        data = DataFrame(
+            {
+                'prediction': [0],
+                'response_time': [0],
+                'timestamp': [input_data['timestamp']],
+                'model_id': [input_data['model_id']],
+                'prediction_confidence': [input_data['prediction_confidence']],
+                'prediction_status': ['Bad'],
+                'comments': [input_data['comment']],
+            }
+        )
 
-        self.info(f"Default prediction formatted: {data.size} rows", metadata)
+        self.info(f'Default prediction formatted: {data.size} rows', metadata)
         return data.to_dict()
 
-    @activity.defn(name="get_last_timestamp")
+    @activity.defn(name='get_last_timestamp')
     async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
         """
         Extract the most recent timestamp from prediction data.
@@ -517,24 +517,22 @@ class Gates(BaseActivity):
         """
         metadata = input_data['metadata']
 
-        self.info("Getting last timestamp...", metadata)
+        self.info('Getting last timestamp...', metadata)
 
         data = DataFrame(input_data['data'])
 
-        self.debug(f"Input data: {data.head(5).to_string()}", metadata)
+        self.debug(f'Input data: {data.head(5).to_string()}', metadata)
 
         if data.empty:
             return now().strftime(DATETIME_FORMAT_WITH_TZ)
 
-        max_timestamp = max(
-            data['timestamp'].values.tolist())
+        max_timestamp = max(data['timestamp'].values.tolist())
 
-        self.info(
-            f"Last timestamp: {max_timestamp}", metadata)
+        self.info(f'Last timestamp: {max_timestamp}', metadata)
 
         return max_timestamp
 
-    @activity.defn(name="write_metrics")
+    @activity.defn(name='write_metrics')
     async def write_metrics(self, input_data: dict[str, Any]):
         """
         Write prediction performance metrics to Prometheus monitoring system.
@@ -562,26 +560,24 @@ class Gates(BaseActivity):
         prediction_confidence = prediction['prediction_confidence'].values[0]
         response_time = prediction['response_time'].values[0]
 
-        self.info(
-            f"Writing metrics for model {metadata['model_name']}", metadata)
+        self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
 
         metrics.PREDICTIONS_WRITTEN_COUNT.labels(
             pod_id=self.pod_id,
             model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
+            pipeline_name=metadata['workflow_name'],
         ).inc()
 
         metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
             pod_id=self.pod_id,
             model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
+            pipeline_name=metadata['workflow_name'],
         ).set(prediction_confidence)
 
         metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
             pod_id=self.pod_id,
             model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name']
+            pipeline_name=metadata['workflow_name'],
         ).observe(response_time)
 
-        self.info(
-            f"Metrics written for model {metadata['model_name']}", metadata)
+        self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

model_manager/activities/mlflow.py

@@ -1,20 +1,19 @@
 from temporalio import activity, workflow
 
-
 with workflow.unsafe.imports_passed_through():
-    from datetime import datetime
-    from pandas import Timestamp, to_datetime
-    from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
-    from sientia_do.temporal.activities.base import BaseActivity
+    import traceback
+    from typing import Any
+
+    import numpy as np
+    from pandas import DataFrame, to_datetime
+    from sientia_do.formatters import create_sample_dict
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
     from sientia_do.observability.logger import Logger
-    from sientia_do.formatters import create_sample_dict
+    from sientia_do.temporal.activities.base import BaseActivity
+    from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
     from model_manager.utils.repository.model_repository import MLFlowRepository
-    from typing import Any
-    import numpy as np
-    from pandas import DataFrame
-    import traceback
 
 
 class MLFlow(BaseActivity):
@@ -36,8 +35,15 @@ class MLFlow(BaseActivity):
         model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
     """
 
-    def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
-                 mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
+    def __init__(
+        self,
+        mlflow_host: str,
+        mlflow_port: int,
+        mlflow_username: str,
+        mlflow_password: str,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
         """
         Initialize MLFlow activities with server configuration.
 
@@ -52,18 +58,17 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If MLFlowRepository initialization fails
         """
-        BaseActivity.__init__(
-            self, logger, notification_handler, set_error_counter=True)
+        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
         self.mlflow_host = mlflow_host
         self.mlflow_port = mlflow_port
         self.mlflow_username = mlflow_username
         self.mlflow_password = mlflow_password
 
         self.model_monitoring_repository = MLFlowRepository(
-            f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger
+            f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
         )
 
-    @activity.defn(name="request_transform")
+    @activity.defn(name='request_transform')
     async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Transform input data using MLFlow models.
@@ -99,7 +104,7 @@ class MLFlow(BaseActivity):
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug("Raw input data:", metadata)
+        self.debug('Raw input data:', metadata)
         self.debug(data.head(5).to_string(), metadata)
 
         # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
@@ -108,14 +113,12 @@ class MLFlow(BaseActivity):
         )
 
         # Pivot data for model input format
-        data = data.pivot(
-            index='timestamp', columns='variable',
-            values='value')
+        data = data.pivot(index='timestamp', columns='variable', values='value')
         data.fillna(np.nan, inplace=True)
         # data.reset_index(inplace=True)
         data.columns.name = None
 
-        self.debug("Processed input data:", metadata)
+        self.debug('Processed input data:', metadata)
         self.debug(data.head(5).to_string(), metadata)
 
         # Request transformation from MLFlow model
@@ -124,16 +127,20 @@ class MLFlow(BaseActivity):
         )
 
         self.debug(
-            f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
         self.debug(
-            f"Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data transformed successfully", metadata)
+        self.info('Data transformed successfully', metadata)
 
         return response_data
 
-    @activity.defn(name="request_predict")
+    @activity.defn(name='request_predict')
     async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Execute predictions using MLFlow models.
@@ -169,14 +176,15 @@ class MLFlow(BaseActivity):
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug(f"Input data for: \n {data.head(5).to_string()}", metadata)
+        self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
 
         # Convert numpy.nan to None for model compatibility
         data.replace(np.nan, None, inplace=True)
 
         data['timestamp'] = data.index
         data['timestamp'] = to_datetime(
-            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
+        ).dt.strftime(DATETIME_FORMAT)
 
         # Request prediction from MLFlow model
         response_data = self.model_monitoring_repository.predict(
@@ -184,13 +192,15 @@ class MLFlow(BaseActivity):
         )
 
         self.debug(
-            f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data predicted successfully", metadata)
+        self.info('Data predicted successfully', metadata)
 
         return response_data
 
-    @activity.defn(name="retrain_model")
+    @activity.defn(name='retrain_model')
     async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Retrain MLFlow models with updated training data.
@@ -234,8 +244,7 @@ class MLFlow(BaseActivity):
         data.drop(columns=['model_id'], inplace=True, errors='ignore')
         data.drop(columns=['created_at'], inplace=True, errors='ignore')
 
-        data = data.pivot(index='timestamp', columns='variable',
-                          values='value')
+        data = data.pivot(index='timestamp', columns='variable', values='value')
         data.sort_index(inplace=True)
         data.reset_index(inplace=True)
 
@@ -244,15 +253,10 @@ class MLFlow(BaseActivity):
 
         try:
             retrain_output, experiment = self.model_monitoring_repository.retrain_model(
-                data=data,
-                model_name=model_name
+                data=data, model_name=model_name
             )
 
-            return {
-                'status': retrain_output,
-                'timestamp': timestamp,
-                'experiment': experiment
-            }
+            return {'status': retrain_output, 'timestamp': timestamp, 'experiment': experiment}
         except Exception as e:
             trace = traceback.format_exc()
             self.send_notification(
@@ -261,12 +265,12 @@ class MLFlow(BaseActivity):
                 message=f'Error retraining model {model_name}: {e}',
                 block='retrain_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             self.error(trace, metadata=metadata)
             raise e
 
-    @activity.defn(name="update_production_model")
+    @activity.defn(name='update_production_model')
     async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Update production model with newly trained model version.
@@ -311,12 +315,12 @@ class MLFlow(BaseActivity):
         status = input_data['status']
 
         self.info(
-            f'Updating production model {model_name} from experiment {experiment}...', metadata)
+            f'Updating production model {model_name} from experiment {experiment}...', metadata
+        )
 
         try:
             response = self.model_monitoring_repository.update_production_model(
-                experiment=experiment,
-                model_name=model_name
+                experiment=experiment, model_name=model_name
             )
 
             report = DataFrame([response])
@@ -325,8 +329,7 @@ class MLFlow(BaseActivity):
             report['timestamp'] = timestamp
             report['status'] = status
 
-            self.info(
-                f'Production model {model_name} updated successfully', metadata)
+            self.info(f'Production model {model_name} updated successfully', metadata)
             return report.to_dict()
 
         except Exception as e:
@@ -337,7 +340,7 @@ class MLFlow(BaseActivity):
                 message=f'Error updating production model {model_name}: {e}',
                 block='update_production_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             self.error(trace, metadata=metadata)
             raise e

model_manager/metrics.py

@@ -22,36 +22,36 @@ Metric Labels:
 - pipeline_name: Name of the prediction pipeline
 """
 
-from prometheus_client import Gauge, Counter, Histogram
+from prometheus_client import Counter, Gauge, Histogram
 
 # Application health metric
 APP_UP = Gauge(
-    "app_up",
-    "Indicates if the application is running (1) or shutting down (0)",
-    ["pod_id"],
+    'app_up',
+    'Indicates if the application is running (1) or shutting down (0)',
+    ['pod_id'],
 )
 
 # Core labels used across multiple metrics
-CORE_LABELS = ["pod_id", "model_name", "pipeline_name"]
+CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
 
 # Prediction operation metrics
 PREDICTIONS_WRITTEN_COUNT = Counter(
-    "model_manager_predictions_written_count",
-    "Number of predictions written to the database table predictions",
+    'model_manager_predictions_written_count',
+    'Number of predictions written to the database table predictions',
     CORE_LABELS,
 )
 
 # Prediction quality metrics
 PREDICTION_CONFIDENCE_MONITOR = Gauge(
-    "model_manager_prediction_confidence_monitor",
-    "Current confidence of each prediction",
+    'model_manager_prediction_confidence_monitor',
+    'Current confidence of each prediction',
     CORE_LABELS,
 )
 
 # Performance monitoring metrics
 PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
-    "model_manager_prediction_response_time_monitor",
-    "Current response time of each prediction",
+    'model_manager_prediction_response_time_monitor',
+    'Current response time of each prediction',
     CORE_LABELS,
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )

model_manager/utils/connectors_config.py

@@ -1,9 +1,8 @@
 from os import getenv
-import json
-from typing import Dict, Any
+from typing import Any
 
 
-def build_postgres_config() -> Dict[str, Any]:
+def build_postgres_config() -> dict[str, Any]:
     """
     Build PostgreSQL database configuration from environment variables.
 
@@ -30,11 +29,11 @@ def build_postgres_config() -> Dict[str, Any]:
         'password': getenv('POSTGRES_PASSWORD', 'sientia'),
         'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
         'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
-        'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20'))
+        'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
     }
 
 
-def build_mlflow_config() -> Dict[str, Any]:
+def build_mlflow_config() -> dict[str, Any]:
     """
     Build MLFlow server configuration from environment variables.
 
@@ -55,11 +54,11 @@ def build_mlflow_config() -> Dict[str, Any]:
         'host': getenv('MLFLOW_HOST', 'http://localhost'),
         'port': int(getenv('MLFLOW_PORT', '5080')),
         'username': getenv('MLFLOW_USERNAME', 'aignosi'),
-        'password': getenv('MLFLOW_PASSWORD', 'aignosi')
+        'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
     }
 
 
-def build_mongodb_config() -> Dict[str, Any]:
+def build_mongodb_config() -> dict[str, Any]:
     """
     Build MongoDB configuration from environment variables.
 
@@ -86,5 +85,5 @@ def build_mongodb_config() -> Dict[str, Any]:
     return {
         'connection_string': connection_string,
         'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
-        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
+        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
     }

model_manager/utils/filters/conditional_filters.py

@@ -20,8 +20,7 @@ def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool
               False if none of the specified variables contain null values.
 
     """
-    return not data[
-        data['variable'].isin(config['variables']) & data['value'].isna()].empty
+    return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
 
 
 def filter_empty_data(data: DataFrame, _config: dict) -> bool:

model_manager/utils/filters/mlflow_filters.py

@@ -52,8 +52,11 @@ def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
         bool: True if data should be filtered (too many NaN values), False otherwise
 
     """
-    data = predictions.replace({None: np.nan}).drop(
-        columns=['timestamp'], errors='ignore').infer_objects()
+    data = (
+        predictions.replace({None: np.nan})
+        .drop(columns=['timestamp'], errors='ignore')
+        .infer_objects()
+    )
 
     if data.isna().all().all():
         return True

model_manager/utils/repository/model_repository.py

@@ -10,22 +10,23 @@ requests using the Model Monitoring API functions.
 By Monitoring we mean the evaluation of the performance of models, the generation of reports.
 
 """
-from datetime import datetime
+
 import traceback
-import pandas as pd
-import mlflow
+from datetime import datetime
 from os import makedirs, path, remove
+
+import mlflow
+import pandas as pd
 from sientia.ModelServing import ModelServing
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 from sientia_do.observability.logger import Logger
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
-class MLFlowRepository():
+class MLFlowRepository:
     def __init__(self, host, username, password, logger: Logger):
-
-        self.model_serving = ModelServing(tracking_uri=host,
-                                          username=username, password=password,
-                                          logger=logger)
+        self.model_serving = ModelServing(
+            tracking_uri=host, username=username, password=password, logger=logger
+        )
         self.logger = logger
 
     def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
@@ -40,33 +41,32 @@ class MLFlowRepository():
         # Get type of first element of index
         index_type = type(index[0])
 
-        self.logger.custom_info(f"Index type: {index_type}", metadata)
+        self.logger.custom_info(f'Index type: {index_type}', metadata)
 
-        message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
+        message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}'
 
         # Check if all in index are of the same type
         if not all(isinstance(i, index_type) for i in index):
-            raise ValueError(
-                f"{message}")
+            raise ValueError(f'{message}')
 
         # Check type and converts to DATETIME_FORMAT_WITH_TZ
-        if index_type == str:
+        if index_type is str:
             # Validate format of string and return error if not valid
             try:
                 pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
-            except ValueError:
-                raise ValueError(
-                    f"{message}")
+            except ValueError as e:
+                raise ValueError(f'{message}') from e
 
         elif index_type == datetime or index_type == pd.Timestamp:
             data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
         else:
-            raise ValueError(
-                f"{message}")
+            raise ValueError(f'{message}')
 
         return data
 
-    def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
+    def transform(
+        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
+    ) -> dict:
         """
         Transform data using a model.
 
@@ -81,41 +81,42 @@ class MLFlowRepository():
 
         try:
             self.logger.custom_debug(
-                f"Data received for model transformation: {data.to_csv()}", metadata)
+                f'Data received for model transformation: {data.to_csv()}', metadata
+            )
 
             model_retention = model_config.get('retention_minutes', 0)
             flavor = model_config.get('transform_flavor', 'sklearn')
             compressed = model_config.get('is_compressed', False)
             retention_target = model_config.get('retention_target', 'model')
-            transform_keyword = model_config.get(
-                'transform_function_keyword', 'predict')
+            transform_keyword = model_config.get('transform_function_keyword', 'predict')
 
             transformed_data = self.model_serving.get_cached_transform(
-                model_name, data, model_retention, flavor,
-                compressed, retention_target, transform_keyword
+                model_name,
+                data,
+                model_retention,
+                flavor,
+                compressed,
+                retention_target,
+                transform_keyword,
             )
 
             self.logger.custom_debug(
-                f"Data received from model transformation: {transformed_data.to_csv()}", metadata)
+                f'Data received from model transformation: {transformed_data.to_csv()}', metadata
+            )
 
-            transformed_data = self.detect_and_parse_datetime_index(
-                transformed_data, metadata)
+            transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata)
 
-            return {
-                'success': True,
-                'content': transformed_data.to_dict()
-            }
+            return {'success': True, 'content': transformed_data.to_dict()}
 
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             return {
                 'success': False,
-                'content': {
-                    'message': str(e),
-                    'traceback': traceback.format_exc()
-                }
+                'content': {'message': str(e), 'traceback': traceback.format_exc()},
             }
 
-    def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict) -> dict:
+    def predict(
+        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
+    ) -> dict:
         """
         Predict data using a model.
 
@@ -137,31 +138,26 @@ class MLFlowRepository():
             start_time = datetime.now()
 
             self.logger.custom_debug(
-                f"Data received for model prediction: {data.to_csv()}", metadata)
+                f'Data received for model prediction: {data.to_csv()}', metadata
+            )
             data = self.model_serving.get_cached_predict(
-                model_name, data, model_retention, flavor,
-                compressed, retention_target
+                model_name, data, model_retention, flavor, compressed, retention_target
             )
 
             end_time = datetime.now()
             data = pd.DataFrame(data, columns=['prediction'])
             self.logger.custom_debug(
-                f"Data received from model prediction: {data.to_csv()}", metadata)
+                f'Data received from model prediction: {data.to_csv()}', metadata
+            )
             data.index = input_index
             data['response_time'] = (end_time - start_time).total_seconds()
 
-            return {
-                'success': True,
-                'content': data.to_dict()
-            }
+            return {'success': True, 'content': data.to_dict()}
 
-        except Exception as e:
+        except Exception as e:  # noqa: BLE001
             return {
                 'success': False,
-                'content': {
-                    'message': str(e),
-                    'traceback': traceback.format_exc()
-                }
+                'content': {'message': str(e), 'traceback': traceback.format_exc()},
             }
 
     def get_experiment_by_run_id(self, run_id: str) -> dict:
@@ -190,10 +186,9 @@ class MLFlowRepository():
         Returns:
             str: The next run name in format 'model_name-run_number'
         """
-        runs = mlflow.search_runs(
-            experiment_names=[model_name], order_by=["start_time desc"])
+        runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc'])
         next_run_number = len(runs) + 1
-        return f"{model_name}-{next_run_number}"
+        return f'{model_name}-{next_run_number}'
 
     def create_model_experiment(self, model_name: str, data: pd.DataFrame) -> tuple:
         """
@@ -217,14 +212,10 @@ class MLFlowRepository():
                 - experiment: MLFlow experiment name
         """
         # load predictor model
-        predictor_uri = f"models:/{model_name}/production"
+        predictor_uri = f'models:/{model_name}/production'
         # load transform model
-        latest_production_id = self.model_serving.get_model_run_id(
-            model_name, stage="Production"
-        )
-        transform_uri = self.model_serving.get_model_uri(
-            latest_production_id, prediction=False
-        )
+        latest_production_id = self.model_serving.get_model_run_id(model_name, stage='Production')
+        transform_uri = self.model_serving.get_model_uri(latest_production_id, prediction=False)
         # load
         data_model = mlflow.sklearn.load_model(transform_uri)
         prediction_model = mlflow.sklearn.load_model(predictor_uri)
@@ -233,20 +224,16 @@ class MLFlowRepository():
 
         target_name = data_model.target_variable
         y = data[target_name]
-        treated_data = pd.merge(
-            treated_data, y, left_index=True, right_index=True)
+        treated_data = pd.merge(treated_data, y, left_index=True, right_index=True)
         prediction_model = prediction_model.fit(treated_data)
         experiment = self.get_experiment_by_run_id(latest_production_id)
         mlflow.set_experiment(experiment)
 
         return prediction_model, data_model, experiment
 
-    def perform_model_retrain(self,
-                              prediction_model,
-                              data_model,
-                              experiment: str,
-                              model_name: str,
-                              data: pd.DataFrame):
+    def perform_model_retrain(
+        self, prediction_model, data_model, experiment: str, model_name: str, data: pd.DataFrame
+    ):
         """
         Execute the complete model retraining process in MLFlow.
 
@@ -271,7 +258,7 @@ class MLFlowRepository():
         """
         pred_model_atributes = vars(prediction_model)  # load class attributes
         data_model_atributes = vars(data_model)  # load class attributes
-        experiment_description = f"Retrain model {model_name} with new data"
+        experiment_description = f'Retrain model {model_name} with new data'
         current_run_name = self.get_next_run_name(experiment)
         with mlflow.start_run(
             run_name=current_run_name, description=experiment_description
@@ -279,32 +266,32 @@ class MLFlowRepository():
             # update transfomation model
             # fixed parameters
             for name_atribute, val_atribute in pred_model_atributes.items():
-                if name_atribute != "model":
+                if name_atribute != 'model':
                     mlflow.log_param(name_atribute, val_atribute)
             # update prediction model
             for name_atribute, val_atribute in data_model_atributes.items():
-                if name_atribute != "model":
+                if name_atribute != 'model':
                     mlflow.log_param(name_atribute, val_atribute)
             # dynamic parameters, including model itself
-            mlflow.sklearn.log_model(data_model, "data_model")
+            mlflow.sklearn.log_model(data_model, 'data_model')
 
-            makedirs("temp", exist_ok=True)
+            makedirs('temp', exist_ok=True)
 
-            file_path = f"temp/raw_data_{model_name}.csv"
+            file_path = f'temp/raw_data_{model_name}.csv'
             data.to_csv(file_path, index=True)
 
             # log the data raw
             mlflow.log_artifact(file_path)
 
             # dynamic parameters, including model itself
-            mlflow.sklearn.log_model(prediction_model, "prediction_model")
-            mlflow.log_param("retrain", True)
+            mlflow.sklearn.log_model(prediction_model, 'prediction_model')
+            mlflow.log_param('retrain', True)
 
             # clear temp file
             if path.exists(file_path):
                 remove(file_path)
 
-            return "Model retrained successfully", experiment
+            return 'Model retrained successfully', experiment
 
     def retrain_model(self, data: pd.DataFrame, model_name: str) -> tuple:
         """
@@ -325,10 +312,10 @@ class MLFlowRepository():
                 - status_message (str): Retraining operation status
                 - experiment_name (str): MLFlow experiment identifier
         """
-        prediction_model, data_model, experiment = self.create_model_experiment(
-            model_name, data)
+        prediction_model, data_model, experiment = self.create_model_experiment(model_name, data)
         retrain_result = self.perform_model_retrain(
-            prediction_model, data_model, experiment, model_name, data)
+            prediction_model, data_model, experiment, model_name, data
+        )
         return retrain_result
 
     def get_experiment(self, experiment_name: str) -> int:
@@ -374,22 +361,21 @@ class MLFlowRepository():
         """
         runs = mlflow.search_runs(
             experiment_ids=[experiment_id],
-            filter_string="",  # Sem filtro no MLflow ainda
-            output_format="pandas"
+            filter_string='',  # Sem filtro no MLflow ainda
+            output_format='pandas',
         )
 
         if not isinstance(runs, pd.DataFrame):
             raise ValueError('Runs is not a pandas DataFrame')
 
         # Filtrar apenas as runs onde params.retrain == True
-        filtered_runs = runs[runs["params.retrain"] == 'True']
+        filtered_runs = runs[runs['params.retrain'] == 'True']
 
         # Converter a coluna 'end_time' para datetime
         filtered_runs['end_time'] = pd.to_datetime(filtered_runs['end_time'])
 
         # Ordenar o DataFrame de forma descendente pela coluna 'end_time'
-        filtered_runs = filtered_runs.sort_values(
-            by='end_time', ascending=False)
+        filtered_runs = filtered_runs.sort_values(by='end_time', ascending=False)
 
         # Pegar a última run_id do DataFrame filtrado e ordenado
         latest_run_id = filtered_runs.iloc[0]['run_id']
@@ -423,16 +409,14 @@ class MLFlowRepository():
         # Registrar o modelo
         # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
         # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
-        mlflow.register_model(
-            f"runs:/{run_id}/prediction_model", model_name)
+        mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
 
         # Colocar a versão do modelo em produção
         # Depois de registrar o modelo, precisamos pegar a versão mais recente do modelo e movê-lo para o estágio 'Production'
         client = mlflow.tracking.MlflowClient()
 
         # Obter a versão mais recente registrada do modelo
-        model_versions = client.get_registered_model(
-            model_name).latest_versions
+        model_versions = client.get_registered_model(model_name).latest_versions
 
         if not isinstance(model_versions, list):
             raise ValueError('Model versions is not a list')
@@ -441,17 +425,10 @@ class MLFlowRepository():
 
         # Mover a versão mais recente do modelo para o estágio de 'Production'
         client.transition_model_version_stage(
-            name=model_name,
-            version=max_version,
-            stage="Production",
-            archive_existing_versions=True
+            name=model_name, version=max_version, stage='Production', archive_existing_versions=True
         )
 
-        return {
-            'model_name': model_name,
-            'version': max_version,
-            'mlflow_run_id': run_id
-        }
+        return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
 
     def update_production_model(self, experiment: str, model_name: str) -> dict:
         """

model_manager/worker/worker.py

@@ -25,32 +25,35 @@ Environment Variables:
 - PROJECT_NAME: Project name for notifications (default: model_manager)
 """
 
-from temporalio import workflow, client
-from temporalio.worker import Worker, PollerBehaviorAutoscaling
-from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
+from temporalio import client, workflow
+from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
+from temporalio.worker import PollerBehaviorAutoscaling, Worker
 
 with workflow.unsafe.imports_passed_through():
+    import asyncio
     import os
     import sys
-    import asyncio
-    from model_manager.workflows.minimal_retrain import MinimalRetrain
-    from model_manager.workflows.predictions_batch import PredictionsBatch
-    from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
-    from model_manager.workflows.sub_workflows.format_and_export_prediction import \
-        FormatAndExportPrediction
-    from model_manager.activities.activities import Activities
-    from model_manager.utils.connectors_config import (
-        build_postgres_config,
-        build_mlflow_config,
-        build_mongodb_config
-    )
+
+    from prometheus_client import start_http_server
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import get_logger
+
     from model_manager import metrics
-    from prometheus_client import start_http_server
+    from model_manager.activities.activities import Activities
+    from model_manager.utils.connectors_config import (
+        build_mlflow_config,
+        build_mongodb_config,
+        build_postgres_config,
+    )
+    from model_manager.workflows.minimal_retrain import MinimalRetrain
+    from model_manager.workflows.predictions_batch import PredictionsBatch
+    from model_manager.workflows.sub_workflows.format_and_export_prediction import (
+        FormatAndExportPrediction,
+    )
+    from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
 
 POD_ID = os.getenv('POD_ID')
-SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
+SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
 
 
 async def main():
@@ -85,7 +88,7 @@ async def main():
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
 
-    logger.custom_info("Starting prometheus client...", metadata)
+    logger.custom_info('Starting prometheus client...', metadata)
     start_prometheus_server()
 
     logger.custom_info('Starting Notification Handler...', metadata)
@@ -95,7 +98,7 @@ async def main():
         connection_string=mongo_config['connection_string'],
         database=mongo_config['database_name'],
         logger=logger,
-        project_name=os.getenv('PROJECT_NAME', 'model-manager')
+        project_name=os.getenv('PROJECT_NAME', 'model-manager'),
     )
 
     logger.custom_info('Starting Activities...', metadata)
@@ -104,16 +107,14 @@ async def main():
         postgres_config=build_postgres_config(),
         mlflow_config=build_mlflow_config(),
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
-    logger.custom_info(
-        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
+    logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
 
     new_runtime = Runtime(
         telemetry=TelemetryConfig(
-            metrics=PrometheusConfig(
-                bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
+            metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
         )
     )
 
@@ -122,7 +123,7 @@ async def main():
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
-        runtime=new_runtime
+        runtime=new_runtime,
     )
 
     logger.custom_info('Starting Workers...', metadata)
@@ -136,20 +137,19 @@ async def main():
                 activities.load_custom_query,
                 activities.retrain_model,
                 activities.update_production_model,
-                activities.export_data_to_postgres
+                activities.export_data_to_postgres,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,
             max_concurrent_local_activities=50,
             max_cached_workflows=200,
             workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
+            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
         ),
         Worker(
             temporal_client,
             task_queue='predictions_batch-queue',
-            workflows=[PredictionsBatch, PredictionProcess,
-                       FormatAndExportPrediction],
+            workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
             activities=[
                 # MLFlow
                 activities.request_predict,
@@ -165,15 +165,15 @@ async def main():
                 activities.load_custom_query,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres,
-                activities.write_metrics
+                activities.write_metrics,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,
             max_concurrent_local_activities=50,
             max_cached_workflows=200,
             workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
-        )
+            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
+        ),
     ]
 
     handlers = []
@@ -186,8 +186,8 @@ async def main():
         # This will run the workers and wait for them to complete.
         # If an exception occurs in any of the worker handlers, it will be propagated here.
         await asyncio.gather(*handlers)
-    except BaseException as e:  # NOSONAR
-        logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
+    except BaseException as e:  # noqa: BLE001
+        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
     finally:
         if notification_handler:
             notification_handler.shutdown()
@@ -216,12 +216,12 @@ def start_prometheus_server():
         SystemExit: If the metrics server fails to start
     """
     try:
-        port = int(os.getenv("HTTP_METRICS_PORT", 9090))
+        port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f"Prometheus server started on port {port}.")
+        print(f'Prometheus server started on port {port}.')
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
-    except Exception as e:
-        print(f"Failed to start Prometheus server: {e}")
+    except Exception as e:  # noqa: BLE001
+        print(f'Failed to start Prometheus server: {e}')
         os._exit(1)
 
 

model_manager/workflows/minimal_retrain.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from model_manager.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from model_manager.activities.activities import Activities
 
 
-@workflow.defn(name="minimal_retrain")
-class MinimalRetrain():
+@workflow.defn(name='minimal_retrain')
+class MinimalRetrain:
     """
     Automated model retraining workflow for the Model Manager system.
 
@@ -63,7 +65,7 @@ class MinimalRetrain():
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'workflow_name': 'minimal_retrain'
+                'workflow_name': 'minimal_retrain',
             }
         }
 
@@ -74,21 +76,17 @@ class MinimalRetrain():
             {
                 **metadata,
                 'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
+                'datetime_columns': input_data.get('datetime_columns', []),
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         experiment_response = await workflow.execute_activity_method(
             Activities.retrain_model,
-            {
-                **metadata,
-                'data': data,
-                'model_name': model_name
-            },
+            {**metadata, 'data': data, 'model_name': model_name},
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         report = await workflow.execute_activity_method(
@@ -97,10 +95,10 @@ class MinimalRetrain():
                 **metadata,
                 'model_name': model_name,
                 'model_id': input_data['model_id'],
-                **experiment_response
+                **experiment_response,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         await workflow.execute_activity_method(
@@ -109,8 +107,8 @@ class MinimalRetrain():
                 **metadata,
                 'data': report,
                 'schema': input_data['schema'],
-                'table_name': input_data['table_name']
+                'table_name': input_data['table_name'],
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )

model_manager/workflows/predictions_batch.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from model_manager.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from model_manager.activities.activities import Activities
 
 
-@workflow.defn(name="predictions_batch")
-class PredictionsBatch():
+@workflow.defn(name='predictions_batch')
+class PredictionsBatch:
     """
     Main batch prediction workflow for the Model Manager system.
 
@@ -74,7 +76,7 @@ class PredictionsBatch():
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'workflow_name': 'predictions_batch'
+                'workflow_name': 'predictions_batch',
             }
         }
 
@@ -84,10 +86,10 @@ class PredictionsBatch():
             {
                 **metadata,
                 'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
+                'datetime_columns': input_data.get('datetime_columns', []),
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=300)
+            start_to_close_timeout=timedelta(seconds=300),
         )
 
         # Prepare input for prediction_process workflow
@@ -98,28 +100,17 @@ class PredictionsBatch():
             'table_name': input_data['table_name'],
             'model_id': input_data['model_id'],
             'model_name': input_data['model_name'],
-            'input_filters': input_data.get('input_filters', {
-                'EMPTY_DATA': {
-                    'POLICY': 'STOP'
-                }
-            }),
-            'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
-                'API_ERROR': {
-                    'POLICY': 'STOP'
-                }
-            }),
-            'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
-                'API_ERROR': {
-                    'POLICY': 'STOP'
-                }
-            }),
+            'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
+            'mlflow_transform_filters': input_data.get(
+                'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
+            ),
+            'mlflow_predict_filters': input_data.get(
+                'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
+            ),
             'model_config': input_data.get('model_config', {}),
             'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
-
-            'prediction_store_policy': input_data.get(
-                'prediction_store_policy', 'lts:1')
+            'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
         }
 
         # Execute prediction process workflow
-        await workflow.execute_child_workflow(
-            'prediction_process', prediction_input)
+        await workflow.execute_child_workflow('prediction_process', prediction_input)

model_manager/workflows/sub_workflows/format_and_export_prediction.py

@@ -1,15 +1,17 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from model_manager.activities.activities import Activities
-    from typing import Any
     from datetime import timedelta
-    from sientia_do.temporal.policies import retry_policy
+    from typing import Any
+
     from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+    from sientia_do.temporal.policies import retry_policy
+
+    from model_manager.activities.activities import Activities
 
 
-@workflow.defn(name="format_and_export_prediction")
-class FormatAndExportPrediction():
+@workflow.defn(name='format_and_export_prediction')
+class FormatAndExportPrediction:
     """
     Data formatting and export workflow for prediction results.
 
@@ -75,10 +77,10 @@ class FormatAndExportPrediction():
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
                     'prediction_confidence': prediction_confidence,
-                    'prediction_store_policy': input_data['prediction_store_policy']
+                    'prediction_store_policy': input_data['prediction_store_policy'],
                 },
                 retry_policy=retry_policy,
-                start_to_close_timeout=timedelta(seconds=60)
+                start_to_close_timeout=timedelta(seconds=60),
             )
 
         else:
@@ -90,10 +92,10 @@ class FormatAndExportPrediction():
                     'timestamp': input_data['timestamp'],
                     'model_id': input_data['model_id'],
                     'prediction_confidence': prediction_confidence,
-                    'comment': input_data['comment']
+                    'comment': input_data['comment'],
                 },
                 retry_policy=retry_policy,
-                start_to_close_timeout=timedelta(seconds=60)
+                start_to_close_timeout=timedelta(seconds=60),
             )
 
         # write to postgres
@@ -104,21 +106,15 @@ class FormatAndExportPrediction():
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
                 'data': prediction,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
+                'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         await workflow.execute_activity_method(
             Activities.write_metrics,
-            {
-                **metadata,
-                'prediction': prediction
-            },
+            {**metadata, 'prediction': prediction},
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )

model_manager/workflows/sub_workflows/prediction_process.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from model_manager.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from model_manager.activities.activities import Activities
 
 
-@workflow.defn(name="prediction_process")
-class PredictionProcess():
+@workflow.defn(name='prediction_process')
+class PredictionProcess:
     """
     Core prediction processing workflow for the Model Manager system.
 
@@ -84,10 +86,7 @@ class PredictionProcess():
         # Get last timestamp for incremental processing
         last_timestamp = await workflow.execute_local_activity_method(
             Activities.get_last_timestamp,
-            {
-                **metadata,
-                'data': data
-            },
+            {**metadata, 'data': data},
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
         )
@@ -97,7 +96,7 @@ class PredictionProcess():
             **metadata,
             'filters': input_data['input_filters'],
             'data': data,
-            'path_priority': input_data['path_priority']
+            'path_priority': input_data['path_priority'],
         }
 
         path_flag, confidence, comment = await workflow.execute_local_activity_method(
@@ -116,12 +115,7 @@ class PredictionProcess():
         # Request MLFlow model transformation
         response_data = await workflow.execute_local_activity_method(
             Activities.request_transform,
-            {
-                **metadata,
-                'data': data,
-                'model_name': model_name,
-                'model_config': model_config
-            },
+            {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=5),
         )
@@ -134,7 +128,7 @@ class PredictionProcess():
                 'filters': input_data['mlflow_transform_filters'],
                 'data': response_data,
                 'type': 'transform',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -155,7 +149,7 @@ class PredictionProcess():
                 'filters': input_data['mlflow_transform_filters'],
                 'data': transformed_data,
                 'type': 'transform',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -172,7 +166,7 @@ class PredictionProcess():
                 **metadata,
                 'data': transformed_data,
                 'model_name': model_name,
-                'model_config': model_config
+                'model_config': model_config,
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=5),
@@ -186,7 +180,7 @@ class PredictionProcess():
                 'filters': input_data['mlflow_predict_filters'],
                 'data': response_data,
                 'type': 'predict',
-                'path_priority': input_data['path_priority']
+                'path_priority': input_data['path_priority'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(minutes=1),
@@ -213,12 +207,19 @@ class PredictionProcess():
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
                 'comment': comment,
-                'prediction_store_policy': input_data['prediction_store_policy']
-            }
+                'prediction_store_policy': input_data['prediction_store_policy'],
+            },
         )
 
-    async def path_flag_handler(self, data: dict, path_flag: str, input_data: dict,
-                                confidence: int, last_timestamp: str, comment: str) -> bool:
+    async def path_flag_handler(
+        self,
+        data: dict,
+        path_flag: str,
+        input_data: dict,
+        confidence: int,
+        last_timestamp: str,
+        comment: str,
+    ) -> bool:
         """
         Handle path decisions based on filter results and confidence levels.
 
@@ -264,7 +265,7 @@ class PredictionProcess():
                     'schema': schema,
                     'table_name': table_name,
                     'model': model_id,
-                    'last_timestamp': last_timestamp
+                    'last_timestamp': last_timestamp,
                 },
                 retry_policy=retry_policy,
                 start_to_close_timeout=timedelta(minutes=1),
@@ -286,8 +287,8 @@ class PredictionProcess():
                     'schema': schema,
                     'table_name': table_name,
                     'comment': comment,
-                    'prediction_store_policy': input_data['prediction_store_policy']
-                }
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                },
             )
             return True
 

tests/laborious/activities/test_activities.py

@@ -1,16 +1,17 @@
+from unittest.mock import ANY, MagicMock, patch
+
 from pytest import mark
-from unittest.mock import patch, MagicMock, ANY
 from sientia_do.temporal.activities.postgres import Postgres
+
 from model_manager.activities.activities import Activities
-from model_manager.activities.mlflow import MLFlow
 from model_manager.activities.gates import Gates
+from model_manager.activities.mlflow import MLFlow
 
 
 @patch('model_manager.activities.activities.Postgres.__init__')
 @patch('model_manager.activities.activities.MLFlow.__init__')
 @patch('model_manager.activities.activities.Gates.__init__')
 def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
-
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -18,15 +19,10 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
-        'host': 'localhost',
-        'port': 5000,
-        'username': 'mlflow',
-        'password': 'mlflow'
-    }
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
 
     logger = MagicMock()
     notification_handler = MagicMock()
@@ -35,7 +31,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     assert isinstance(activities, Activities)
@@ -53,7 +49,7 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
         min_connections=postgres_config['min_connections'],
         max_connections=postgres_config['max_connections'],
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     mock_mlflow_init.assert_called_once_with(
@@ -63,13 +59,11 @@ def test___init__(mock_gates_init, mock_mlflow_init, mock_postgres_init):
         mlflow_username=mlflow_config['username'],
         mlflow_password=mlflow_config['password'],
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     mock_gates_init.assert_called_once_with(
-        ANY,
-        logger=logger,
-        notification_handler=notification_handler
+        ANY, logger=logger, notification_handler=notification_handler
     )
 
 
@@ -84,15 +78,10 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
-        'host': 'localhost',
-        'port': 5000,
-        'username': 'mlflow',
-        'password': 'mlflow'
-    }
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
 
     logger = MagicMock()
     notification_handler = MagicMock()
@@ -101,7 +90,7 @@ async def test_shutdown(_mock_mlflow_init, mock_postgres_init):
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     await activities.shutdown()

tests/laborious/activities/test_gates.py

@@ -1,6 +1,8 @@
-from unittest.mock import MagicMock, ANY, patch
+from unittest.mock import ANY, MagicMock, patch
+
 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
+
 from model_manager.activities.gates import Gates
 
 
@@ -20,11 +22,11 @@ def gates_activity():
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
@@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
         'data': {'value': [1, 2, 3]},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.error.assert_called_once_with(
-        "Filter INVALID_FILTER not found", metadata['metadata']
+        'Filter INVALID_FILTER not found', metadata['metadata']
     )
 
 
@@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
     # Arrange
     mock_input_filter_functions.__contains__.return_value = True
     mock_input_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
-        },
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
         'data': {'value': []},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
+        notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
         message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
-        block="input_gate",
+        block='input_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity):
         **metadata,
         'filters': {},
         'data': {'value': [1, 2, 3]},
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -105,18 +104,16 @@ async def test_input_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
-        },
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
         'data': {'value': []},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == ('STOP', -1, "Input data with bad quality")
+    assert result == ('STOP', -1, 'Input data with bad quality')
     gates_activity.debug.assert_called()
 
 
@@ -125,51 +122,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
 
 
 @mark.asyncio
 @patch('model_manager.activities.gates.mlflow_response_filter_functions')
-async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
-                                                     gates_activity):
+async def test_mlflow_response_gate_filter_exception(
+    mock_mlflow_response_filter_functions, gates_activity
+):
     # Arrange
     mock_mlflow_response_filter_functions.__contains__.return_value = True
     mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
+        notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
         message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
-        block="mlflow_gate",
+        block='mlflow_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -181,14 +176,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
         'filters': {},
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -197,25 +192,20 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'API_ERROR': {'policy': 'STOP'}
-        },
+        'filters': {'API_ERROR': {'policy': 'STOP'}},
         'data': {
             'success': False,
-            'content': {
-                'message': 'API error occurred',
-                'traceback': 'error trace'
-            }
+            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
         },
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == ('STOP', -1, "API error occurred")
+    assert result == ('STOP', -1, 'API error occurred')
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called()
 
@@ -225,58 +215,53 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
         'data': {'value': [1, 2, 3]},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
 
 
 @mark.asyncio
 @patch('model_manager.activities.gates.mlflow_content_filter_functions')
-async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
-                                                    gates_activity):
+async def test_mlflow_content_gate_filter_exception(
+    mock_mlflow_content_filter_functions, gates_activity
+):
     # Arrange
     mock_mlflow_content_filter_functions.__contains__.return_value = True
     mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
-            'API_ERROR': {'POLICY': 'STOP'}
-        },
+        'filters': {'API_ERROR': {'POLICY': 'STOP'}},
         'data': {
             'success': False,
-            'content': {
-                'message': 'API error occurred',
-                'traceback': 'error trace'
-            }
+            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
         },
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
+        notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
         message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
-        block="mlflow_gate",
+        block='mlflow_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -288,14 +273,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
         'filters': {},
         'data': {'value': [1, 2, 3]},
         'type': 'test',
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -304,20 +289,17 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
-            'NAN_VALUES': {'policy': 'STOP', 'config': {}}
-        },
+        'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
         'data': {'value': [None, None, None]},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (
-        'STOP', -1, "Transformed data not passed the content filter")
+    assert result == ('STOP', -1, 'Transformed data not passed the content filter')
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called()
 
@@ -328,7 +310,8 @@ def test_get_prediction_store_policy_invalid_policy(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -341,7 +324,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -354,7 +338,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -367,7 +352,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'erl'
@@ -380,16 +366,12 @@ async def test_format_prediction_no_timestamp(gates_activity):
     input_data = {
         **metadata,
         'data': {
-            'prediction': {
-                '2023-05-26 11:12:27': 1
-            },
-            'response_time': {
-                '2023-05-26 11:12:27': 0.1
-            }
+            'prediction': {'2023-05-26 11:12:27': 1},
+            'response_time': {'2023-05-26 11:12:27': 0.1},
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:1'
+        'prediction_store_policy': 'lts:1',
     }
 
     # Act
@@ -402,7 +384,7 @@ async def test_format_prediction_no_timestamp(gates_activity):
     assert result['model_id'] == {0: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9}
     assert result['prediction_status'] == {0: 'Good'}
-    assert result['comments'] == {0: ""}
+    assert result['comments'] == {0: ''}
 
 
 @mark.asyncio
@@ -420,11 +402,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
                 '2023-05-26 11:12:27': 0.1,
                 '2023-05-26 11:12:28': 0.2,
                 '2023-05-26 11:12:29': 0.3,
-            }
+            },
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'erl:2'
+        'prediction_store_policy': 'erl:2',
     }
 
     # Act
@@ -433,12 +415,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
     # Assert
     assert result['prediction'] == {0: 2, 1: 1}
     assert result['response_time'] == {0: 0.2, 1: 0.1}
-    assert result['timestamp'] == {
-        0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
     assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
     assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
-    assert result['comments'] == {0: "", 1: ""}
+    assert result['comments'] == {0: '', 1: ''}
 
 
 @mark.asyncio
@@ -456,11 +437,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
                 '2023-05-26 11:12:27': 0.1,
                 '2023-05-26 11:12:28': 0.2,
                 '2023-05-26 11:12:29': 0.3,
-            }
+            },
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:2'
+        'prediction_store_policy': 'lts:2',
     }
 
     # Act
@@ -469,12 +450,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
     # Assert
     assert result['prediction'] == {0: 3, 1: 2}
     assert result['response_time'] == {0: 0.3, 1: 0.2}
-    assert result['timestamp'] == {
-        0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
     assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
     assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
-    assert result['comments'] == {0: "", 1: ""}
+    assert result['comments'] == {0: '', 1: ''}
 
 
 @mark.asyncio
@@ -482,22 +462,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {'prediction': [1, 2, 3],
-                 'response_time': [0.1, 0.2, 0.3],
-                 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
+        'data': {
+            'prediction': [1, 2, 3],
+            'response_time': [0.1, 0.2, 0.3],
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+        },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:2'
+        'prediction_store_policy': 'lts:2',
     }
-    gates_activity.get_prediction_store_policy = MagicMock(
-        return_value=('invalid', 1))
+    gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
 
     try:
-        result = await gates_activity.format_prediction(input_data)
+        await gates_activity.format_prediction(input_data)
     except ValueError as e:
-        assert str(e) == "Invalid policy type: invalid"
+        assert str(e) == 'Invalid policy type: invalid'
     else:
-        assert False, "Expected ValueError"
+        raise AssertionError('Expected ValueError')
 
 
 @mark.asyncio
@@ -508,7 +489,7 @@ async def test_format_default_prediction(gates_activity):
         'timestamp': '2023-05-26 11:12:27',
         'model_id': 'test_model',
         'prediction_confidence': 0.1,
-        'comment': 'Test comment'
+        'comment': 'Test comment',
     }
 
     # Act
@@ -528,12 +509,7 @@ async def test_format_default_prediction(gates_activity):
 @mark.asyncio
 async def test_get_last_timestamp_with_data(gates_activity):
     # Arrange
-    input_data = {
-        **metadata,
-        'data': {
-            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
-        }
-    }
+    input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
 
     # Act
     result = await gates_activity.get_last_timestamp(input_data)
@@ -545,10 +521,7 @@ async def test_get_last_timestamp_with_data(gates_activity):
 @mark.asyncio
 async def test_get_last_timestamp_no_data(gates_activity):
     # Arrange
-    input_data = {
-        'data': {},
-        **metadata
-    }
+    input_data = {'data': {}, **metadata}
 
     # Act
     result = await gates_activity.get_last_timestamp(input_data)
@@ -567,30 +540,28 @@ async def test_write_metrics(mock_metrics, gates_activity):
         'prediction': {
             'prediction': [1, 2, 3],
             'prediction_confidence': [0.9, 0.8, 0.7],
-            'response_time': [0.1, 0.2, 0.3]
-        }
+            'response_time': [0.1, 0.2, 0.3],
+        },
     }
     await gates_activity.write_metrics(input_data)
     mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
+        pipeline_name=metadata['metadata']['workflow_name'],
     )
     mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
 
     mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
-    )
-    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(
-        0.9
+        pipeline_name=metadata['metadata']['workflow_name'],
     )
+    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
 
     mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
+        pipeline_name=metadata['metadata']['workflow_name'],
     )
     mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
         0.1

tests/laborious/activities/test_mlflow.py

@@ -1,45 +1,42 @@
-from datetime import datetime
 from unittest.mock import ANY, MagicMock, patch
 
 import numpy as np
-from pandas import DataFrame, Timestamp
-from pytest import fixture, mark, raises
-from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
-from model_manager.activities.mlflow import MLFlow
+from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
+from model_manager.activities.mlflow import MLFlow
 
 
-@patch("model_manager.activities.mlflow.MLFlowRepository")
+@patch('model_manager.activities.mlflow.MLFlowRepository')
 def test___init__(mock_mlflow_repository):
     mlflow = MLFlow(
-        mlflow_host="http://localhost",
+        mlflow_host='http://localhost',
         mlflow_port=5000,
-        mlflow_username="admin",
-        mlflow_password="admin",
+        mlflow_username='admin',
+        mlflow_password='admin',
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
     )
 
-    assert mlflow.mlflow_host == "http://localhost"
+    assert mlflow.mlflow_host == 'http://localhost'
     assert mlflow.mlflow_port == 5000
-    assert mlflow.mlflow_username == "admin"
-    assert mlflow.mlflow_password == "admin"
+    assert mlflow.mlflow_username == 'admin'
+    assert mlflow.mlflow_password == 'admin'
 
-    mock_mlflow_repository.assert_called_once_with(
-        "http://localhost:5000", "admin", "admin", ANY
-    )
+    mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
 
 
 @fixture
-@patch("model_manager.activities.mlflow.MLFlowRepository")
+@patch('model_manager.activities.mlflow.MLFlowRepository')
 def mlflow(mock_mlflow_repository):
     mlflow = MLFlow(
-        mlflow_host="http://localhost:5000",
+        mlflow_host='http://localhost:5000',
         mlflow_port=5000,
-        mlflow_username="admin",
-        mlflow_password="admin",
+        mlflow_username='admin',
+        mlflow_password='admin',
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
     )
 
     mlflow.send_notification = MagicMock()
@@ -48,44 +45,67 @@ def mlflow(mock_mlflow_repository):
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("model_manager.activities.mlflow.DataFrame")
-@patch("model_manager.activities.mlflow.max")
+@patch('model_manager.activities.mlflow.DataFrame')
+@patch('model_manager.activities.mlflow.max')
 async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
         **metadata,
         'data': [
-            {'timestamp': '2024-01-01', 'variable': 'var1',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-01', 'variable': 'var2',
-                'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var1',
-                'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var2',
-                'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var1',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
-            {'timestamp': '2024-01-02', 'variable': 'var2',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
+            {
+                'timestamp': '2024-01-01',
+                'variable': 'var1',
+                'value': 1.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-01',
+                'variable': 'var2',
+                'value': 2.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var1',
+                'value': 3.0,
+                'created_at': '2024-01-02 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var2',
+                'value': 4.0,
+                'created_at': '2024-01-02 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var1',
+                'value': 1.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var2',
+                'value': 1.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
         ],
         'model_name': 'test_model',
-        'model_config': {}
+        'model_config': {},
     }
 
     # Mock the transform response
-    expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
-        '2024-01-01', '2024-01-02']}
+    expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
     mlflow.model_monitoring_repository.transform.return_value = expected_response
 
     mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
@@ -114,30 +134,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
 
 
 @mark.asyncio
-@patch("model_manager.activities.mlflow.DataFrame")
-@patch("model_manager.activities.mlflow.to_datetime")
-@patch("model_manager.activities.mlflow.max")
+@patch('model_manager.activities.mlflow.DataFrame')
+@patch('model_manager.activities.mlflow.to_datetime')
+@patch('model_manager.activities.mlflow.max')
 async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
         **metadata,
         'data': {
-            "variable": {
-                "2024-01-01": "var1",
-                "2024-01-02": "var2",
-                "2024-01-03": "var1",
-                "2024-01-04": "var2"
+            'variable': {
+                '2024-01-01': 'var1',
+                '2024-01-02': 'var2',
+                '2024-01-03': 'var1',
+                '2024-01-04': 'var2',
             },
-            "value": {
-                "2024-01-01": 1.0,
-                "2024-01-02": 2.0,
-                "2024-01-03": 3.0,
-                "2024-01-04": 4.0
-            }
+            'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
         },
         'model_name': 'test_model',
-        'model_config': {}
+        'model_config': {},
     }
 
     # Mock the predict response
@@ -148,9 +163,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
     response_data = await mlflow.request_predict(input_data)
 
     mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.replace.assert_called_once_with(
-        np.nan, None, inplace=True
-    )
+    mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
     mock_dataframe.return_value.__setitem__.assert_any_call(
         'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
     )
@@ -158,9 +171,7 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
     mock_to_datetime.assert_called_once_with(
         mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
     )
-    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
-        DATETIME_FORMAT
-    )
+    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
 
     # Verify the response
     assert response_data == expected_response
@@ -174,28 +185,26 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
 @mark.asyncio
 async def test_retrain_model(mlflow):
     data = {
-        "model_id": [4, 5, 6, 7],
-        "created_at": [1, 2, 3, 4],
-        "timestamp": [1, 1, 2, 2],
-        "variable": ["var1", "var2", "var1", "var2"],
-        "value": [1, 2, 3, 4]
+        'model_id': [4, 5, 6, 7],
+        'created_at': [1, 2, 3, 4],
+        'timestamp': [1, 1, 2, 2],
+        'variable': ['var1', 'var2', 'var1', 'var2'],
+        'value': [1, 2, 3, 4],
     }
 
     mlflow.model_monitoring_repository.retrain_model.return_value = (
-        'Model retrained successfully', 'test')
+        'Model retrained successfully',
+        'test',
+    )
 
-    response = await mlflow.retrain_model({
-        **metadata,
-        'data': data,
-        'model_name': 'test_model'
-    })
+    response = await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
 
     mlflow.model_monitoring_repository.retrain_model.assert_called_once()
 
     assert response == {
-        "status": 'Model retrained successfully',
-        "timestamp": 2,
-        "experiment": 'test'
+        'status': 'Model retrained successfully',
+        'timestamp': 2,
+        'experiment': 'test',
     }
 
 
@@ -206,20 +215,16 @@ async def test_retrain_model_error(mlflow):
     )
 
     data = {
-        "model_id": [4, 5, 6, 7],
-        "created_at": [1, 2, 3, 4],
-        "timestamp": [1, 1, 2, 2],
-        "variable": ["var1", "var2", "var1", "var2"],
-        "value": [1, 2, 3, 4]
+        'model_id': [4, 5, 6, 7],
+        'created_at': [1, 2, 3, 4],
+        'timestamp': [1, 1, 2, 2],
+        'variable': ['var1', 'var2', 'var1', 'var2'],
+        'value': [1, 2, 3, 4],
     }
 
     try:
-        await mlflow.retrain_model({
-            **metadata,
-            'data': data,
-            'model_name': 'test_model'
-        })
-    except Exception as e:
+        await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
+    except Exception as e:  # noqa: BLE001
         assert str(e) == 'Error retraining model'
         mlflow.send_notification.assert_called_once_with(
             metadata=metadata['metadata'],
@@ -227,20 +232,18 @@ async def test_retrain_model_error(mlflow):
             message='Error retraining model test_model: Error retraining model',
             block='retrain_model',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
     else:
-        assert False, "No exception raised"
+        raise AssertionError('No exception raised')
 
 
 @mark.asyncio
 async def test_update_production_model(mlflow):
-    mlflow.model_monitoring_repository.update_production_model.return_value = (
-        {
-            "data1": 1,
-            "data2": 2
-        }
-    )
+    mlflow.model_monitoring_repository.update_production_model.return_value = {
+        'data1': 1,
+        'data2': 2,
+    }
 
     input_data = {
         **metadata,
@@ -248,13 +251,14 @@ async def test_update_production_model(mlflow):
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     response = await mlflow.update_production_model(input_data)
 
     mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
-        experiment='test', model_name='test_model')
+        experiment='test', model_name='test_model'
+    )
 
     assert response == {
         'data1': {0: 1},
@@ -262,7 +266,7 @@ async def test_update_production_model(mlflow):
         'model_id': {0: 1},
         'model_name': {0: 'test_model'},
         'timestamp': {0: 2},
-        'status': {0: 'success'}
+        'status': {0: 'success'},
     }
 
 
@@ -278,12 +282,12 @@ async def test_update_production_model_error(mlflow):
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     try:
         await mlflow.update_production_model(input_data)
-    except Exception as e:
+    except Exception as e:  # noqa: BLE001
         assert str(e) == 'Error updating production model'
         mlflow.send_notification.assert_called_once_with(
             metadata=metadata['metadata'],
@@ -291,7 +295,7 @@ async def test_update_production_model_error(mlflow):
             message='Error updating production model test_model: Error updating production model',
             block='update_production_model',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
     else:
-        assert False, "No exception raised"
+        raise AssertionError('No exception raised')

tests/laborious/utils/filters/test_conditional_filters.py

@@ -1,23 +1,29 @@
 from pandas import DataFrame
 
 from model_manager.utils.filters.conditional_filters import (
+    filter_empty_data,
     filter_specific_variables_null_values,
-    filter_empty_data
 )
 
 
 def test_filter_specific_variables_null_values():
-    assert filter_specific_variables_null_values(
-        DataFrame(
-            {'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
-        config={'variables': ['variable2']}) is False
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
+            config={'variables': ['variable2']},
+        )
+        is False
+    )
 
 
 def test_filter_specific_variables_null_values_with_null_values():
-    assert filter_specific_variables_null_values(
-        DataFrame(
-            {'variable': ['variable1', 'variable2'], 'value': [1, None]}),
-        config={'variables': ['variable2']}) is True
+    assert (
+        filter_specific_variables_null_values(
+            DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, None]}),
+            config={'variables': ['variable2']},
+        )
+        is True
+    )
 
 
 def test_filter_empty_data():
@@ -25,6 +31,7 @@ def test_filter_empty_data():
 
 
 def test_filter_empty_data_with_data():
-    assert filter_empty_data(
-        DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}),
-        {}) is False
+    assert (
+        filter_empty_data(DataFrame({'variable': ['variable1', 'variable2'], 'value': [1, 2]}), {})
+        is False
+    )

tests/laborious/utils/filters/test_mlflow_filters.py

@@ -1,22 +1,23 @@
 from pandas import DataFrame
+
 from model_manager.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
 
 
 def test_api_error_filter_invalid_response():
-    assert api_error_filter(None, {}) == True  # NOSONAR
+    assert api_error_filter(None, {})
 
 
 def test_api_error_filter_valid_response_fail():
-    assert api_error_filter({'success': False}, {}) == True
+    assert api_error_filter({'success': False}, {})
 
 
 def test_api_error_filter_valid_response_success():
-    assert api_error_filter({'success': True}, {}) == False
+    assert not api_error_filter({'success': True}, {})
 
 
 def test_nan_values_filter_all_nan_values():
-    assert nan_values_filter(DataFrame({'variable': [None, None]}), {}) == True
+    assert nan_values_filter(DataFrame({'variable': [None, None]}), {})
 
 
 def test_nan_values_filter_no_nan_values():
-    assert nan_values_filter(DataFrame({'variable': [1, 2]}), {}) == False
+    assert not nan_values_filter(DataFrame({'variable': [1, 2]}), {})

tests/laborious/utils/repository/test_model_repository.py

@@ -1,34 +1,33 @@
+from datetime import UTC, datetime
 from unittest.mock import ANY, MagicMock, call, patch
+
 import numpy as np
-from pandas import DataFrame
 import pytest
-from datetime import datetime, timezone
-from pandas import Timestamp
+from pandas import DataFrame, Timestamp
+
 from model_manager.utils.repository.model_repository import MLFlowRepository
 
 
 @pytest.fixture
 def mlflow_repository():
-    with patch('model_manager.utils.repository.model_repository.ModelServing',
-               autospec=True) as mock_model_serving:
+    with patch(
+        'model_manager.utils.repository.model_repository.ModelServing', autospec=True
+    ) as mock_model_serving:
         mock_instance = mock_model_serving.return_value
         mock_instance.get_transformed_data = MagicMock()
 
         repo = MLFlowRepository(
-            host='http://localhost:5000',
-            username='admin',
-            password='admin',
-            logger=MagicMock()
+            host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
         )
         return repo
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
@@ -38,80 +37,56 @@ class Any:
 
 
 invalid_cases = [
-    (
-        {
-            'value': {
-                '2024-01-01 12:00:00': 1,
-                2024: 2
-            }
-        }
-    ),
-    (
-        {
-            'value': {
-                '2024-01-01': 1,
-                '2024-01-02': 2
-            }
-        }
-    ),
-    (
-        {
-            'value': {
-                Any(): 1,
-                Any(): 2
-            }
-        }
-    )
+    ({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
+    ({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
+    ({'value': {Any(): 1, Any(): 2}}),
 ]
 
 
-@pytest.mark.parametrize("data", invalid_cases)
+@pytest.mark.parametrize('data', invalid_cases)
 def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
-    input_data = DataFrame(
-        data
-    )
+    input_data = DataFrame(data)
 
     with pytest.raises(ValueError) as e:
-        mlflow_repository.detect_and_parse_datetime_index(
-            input_data, metadata['metadata'])
+        mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
 
-        assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S"
+        assert (
+            str(e)
+            == 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S'
+        )
 
 
 valid_cases = [
     (
-        {
-            'value': {
-                '2024-01-01 12:00:00+0000': 1,
-                '2024-01-02 12:00:00+0000': 2
-            }
-        }, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000']
+        {'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
+        ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
     ),
     (
         {
             'value': {
-                datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
-                datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
+                datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
+                datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
             }
-        }, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000']
+        },
+        ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
     ),
     (
         {
             'value': {
-                Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
-                Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
+                Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
+                Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
             }
-        }, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000']
+        },
+        ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'],
     ),
 ]
 
 
-@pytest.mark.parametrize("data,expected", valid_cases)
+@pytest.mark.parametrize('data,expected', valid_cases)
 def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
     input_data = DataFrame(data)
 
-    response = mlflow_repository.detect_and_parse_datetime_index(
-        input_data, metadata['metadata'])
+    response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
 
     assert response.index.tolist() == expected
 
@@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository):
 
     mlflow_repository.detect_and_parse_datetime_index = MagicMock()
 
-    output = mlflow_repository.transform(
-        model_name, data, {}, metadata['metadata'])
+    output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
-        model_name, data, 0, 'sklearn', False, 'model', 'predict')
+        model_name, data, 0, 'sklearn', False, 'model', 'predict'
+    )
 
     mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
-        mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata'])
+        mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']
+    )
 
     assert output == {
         'success': True,
-        'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value
+        'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
     }
 
 
@@ -141,80 +117,48 @@ def test_transform_error(mlflow_repository):
     data = MagicMock()
     model_name = 'model'
 
-    mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
-        'error')
+    mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
 
-    output = mlflow_repository.transform(
-        model_name, data, {}, metadata['metadata'])
+    output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
-        model_name, data, 0, 'sklearn', False, 'model', 'predict')
+        model_name, data, 0, 'sklearn', False, 'model', 'predict'
+    )
 
-    assert output == {
-        'success': False,
-        'content': {
-            'message': 'error',
-            'traceback': ANY
-        }
-    }
+    assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
 
 
 def test_predict_success(mlflow_repository):
-    data = DataFrame({
-        'feat_1': {
-            'index_1': 2,
-            'index_2': 3
-        }
-    })
+    data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
     model_name = 'model'
-    mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
-        [2, 3]
-    )
+    mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
 
-    output = mlflow_repository.predict(
-        model_name, data, {}, metadata['metadata'])
+    output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
-        model_name, data, 0, 'pyfunc', False, 'model')
+        model_name, data, 0, 'pyfunc', False, 'model'
+    )
 
     assert output['success'] is True
     assert output['content'] == {
-        'prediction': {
-            'index_1': 2,
-            'index_2': 3
-        }, 'response_time': {
-            'index_1': ANY,
-            'index_2': ANY
-        }
+        'prediction': {'index_1': 2, 'index_2': 3},
+        'response_time': {'index_1': ANY, 'index_2': ANY},
     }
 
 
 def test_predict_error(mlflow_repository):
-    data = DataFrame({
-        'feat_1': {
-            'index_1': 2,
-            'index_2': 3
-        }
-    })
+    data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
     model_name = 'model'
 
-    mlflow_repository.model_serving.get_cached_predict = MagicMock(
-        side_effect=Exception('error')
-    )
+    mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
 
-    output = mlflow_repository.predict(
-        model_name, data, {}, metadata['metadata'])
+    output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
 
     mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
-        model_name, data, 0, 'pyfunc', False, 'model')
+        model_name, data, 0, 'pyfunc', False, 'model'
+    )
 
-    assert output == {
-        'success': False,
-        'content': {
-            'message': 'error',
-            'traceback': ANY
-        }
-    }
+    assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
 
 
 @patch('model_manager.utils.repository.model_repository.mlflow')
@@ -246,8 +190,7 @@ def test_get_next_run_name(mlflow, mlflow_repository):
 
 @patch('model_manager.utils.repository.model_repository.mlflow')
 def test_get_experiment_success(mlflow, mlflow_repository):
-    mlflow.get_experiment_by_name.return_value = MagicMock(
-        experiment_id='0')
+    mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
 
     output = mlflow_repository.get_experiment('test')
 
@@ -263,23 +206,25 @@ def test_get_experiment_error(mlflow, mlflow_repository):
     except ValueError as e:
         assert str(e) == 'Experiment test not found'
     else:
-        assert False
+        raise AssertionError('Expected exception')
 
 
 @patch('model_manager.utils.repository.model_repository.mlflow')
 def test_get_experiment_last_run(mlflow, mlflow_repository):
-    mlflow.search_runs.return_value = DataFrame({
-        'params.retrain': ['True', 'False', 'True', 'False'],
-        'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
-        'run_id': ['0', '1', '2', '3'],
-    })
+    mlflow.search_runs.return_value = DataFrame(
+        {
+            'params.retrain': ['True', 'False', 'True', 'False'],
+            'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
+            'run_id': ['0', '1', '2', '3'],
+        }
+    )
 
     output = mlflow_repository.get_experiment_last_run(0)
 
     mlflow.search_runs.assert_called_once_with(
         experiment_ids=[0],
-        filter_string="",
-        output_format="pandas",
+        filter_string='',
+        output_format='pandas',
     )
 
     assert output == '2'
@@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository):
     except ValueError as e:
         assert str(e) == 'Runs is not a pandas DataFrame'
     else:
-        assert False
+        raise AssertionError('Expected exception')
 
 
 @patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
 @patch('model_manager.utils.repository.model_repository.mlflow.set_experiment')
 def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
-
-    mlflow_repository.model_serving.get_model_run_id = MagicMock(
-        return_value='0')
-    mlflow_repository.model_serving.get_model_uri = MagicMock(
-        return_value='test')
+    mlflow_repository.model_serving.get_model_run_id = MagicMock(return_value='0')
+    mlflow_repository.model_serving.get_model_uri = MagicMock(return_value='test')
     mlflow_repository.get_experiment_by_run_id = MagicMock()
 
     data_model_mock = MagicMock()
@@ -313,29 +255,30 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
     sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
 
     data_model_mock.fit.return_value = data_model_mock
-    data_model_mock.predict.return_value = DataFrame({
-        'x': [10, 20, 30],
-    })
+    data_model_mock.predict.return_value = DataFrame(
+        {
+            'x': [10, 20, 30],
+        }
+    )
     data_model_mock.target_variable = 'y'
 
     prediction_model_mock.fit.return_value = prediction_model_mock
 
-    data = DataFrame({
-        'x': [1, 2, 3],
-        'y': [4, 5, 6]
-    })
+    data = DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
 
     output = mlflow_repository.create_model_experiment('test', data)
 
     mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
-        'test', stage='Production')
-    mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
-        '0', prediction=False)
+        'test', stage='Production'
+    )
+    mlflow_repository.model_serving.get_model_uri.assert_called_once_with('0', prediction=False)
 
-    sklearn.load_model.assert_has_calls([
-        call(mlflow_repository.model_serving.get_model_uri.return_value),
-        call("models:/test/production"),
-    ])
+    sklearn.load_model.assert_has_calls(
+        [
+            call(mlflow_repository.model_serving.get_model_uri.return_value),
+            call('models:/test/production'),
+        ]
+    )
     assert sklearn.load_model.call_count == 2
 
     data_model_mock.fit.assert_called_once_with(data)
@@ -343,21 +286,23 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
 
     fit_args = prediction_model_mock.fit.call_args[0][0]
     assert fit_args.equals(
-        DataFrame({
-            'x': [10, 20, 30],
-            'y': [4, 5, 6],
-        })
+        DataFrame(
+            {
+                'x': [10, 20, 30],
+                'y': [4, 5, 6],
+            }
+        )
     )
 
     mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
 
-    set_experiment.assert_called_once_with(
-        mlflow_repository.get_experiment_by_run_id.return_value
-    )
+    set_experiment.assert_called_once_with(mlflow_repository.get_experiment_by_run_id.return_value)
 
-    assert output == (prediction_model_mock,
-                      data_model_mock,
-                      mlflow_repository.get_experiment_by_run_id.return_value)
+    assert output == (
+        prediction_model_mock,
+        data_model_mock,
+        mlflow_repository.get_experiment_by_run_id.return_value,
+    )
 
 
 @patch('model_manager.utils.repository.model_repository.mlflow.start_run')
@@ -365,41 +310,43 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
 @patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
 @patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
 def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
-
     prediction_model_mock = MagicMock()
     data_model_mock = MagicMock()
     experiment = 'test'
     model_name = 'test'
     data = MagicMock()
 
-    mlflow_repository.get_next_run_name = MagicMock(
-        return_value='test-1')
+    mlflow_repository.get_next_run_name = MagicMock(return_value='test-1')
     run = MagicMock()
     start_run.__enter__.return_value = run
 
     output = mlflow_repository.perform_model_retrain(
-        prediction_model_mock, data_model_mock, experiment, model_name, data)
+        prediction_model_mock, data_model_mock, experiment, model_name, data
+    )
 
     mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
     start_run.assert_called_once_with(
-        run_name='test-1', description='Retrain model test with new data')
+        run_name='test-1', description='Retrain model test with new data'
+    )
 
-    log_model.assert_has_calls([
-        call(data_model_mock, "data_model"),
-        call(prediction_model_mock, "prediction_model"),
-    ])
+    log_model.assert_has_calls(
+        [
+            call(data_model_mock, 'data_model'),
+            call(prediction_model_mock, 'prediction_model'),
+        ]
+    )
 
-    data.to_csv.assert_called_once_with(
-        "temp/raw_data_test.csv", index=True)
+    data.to_csv.assert_called_once_with('temp/raw_data_test.csv', index=True)
 
-    log_artifact.assert_called_once_with(
-        "temp/raw_data_test.csv")
+    log_artifact.assert_called_once_with('temp/raw_data_test.csv')
 
-    log_param.assert_has_calls([
-        call("retrain", True),
-    ])
+    log_param.assert_has_calls(
+        [
+            call('retrain', True),
+        ]
+    )
 
-    assert output == ("Model retrained successfully", experiment)
+    assert output == ('Model retrained successfully', experiment)
 
 
 def test_retrain_model(mlflow_repository):
@@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository):
     model_name = 'test'
 
     mlflow_repository.create_model_experiment = MagicMock(
-        return_value=('data_model', 'prediction_model', '0'))
+        return_value=('data_model', 'prediction_model', '0')
+    )
 
-    mlflow_repository.perform_model_retrain = MagicMock(
-        return_value='Model retrained successfully')
+    mlflow_repository.perform_model_retrain = MagicMock(return_value='Model retrained successfully')
 
     output = mlflow_repository.retrain_model(data, model_name)
 
-    mlflow_repository.create_model_experiment.assert_called_once_with(
-        model_name, data)
+    mlflow_repository.create_model_experiment.assert_called_once_with(model_name, data)
 
     mlflow_repository.perform_model_retrain.assert_called_once_with(
-        'data_model', 'prediction_model', '0', model_name, data)
+        'data_model', 'prediction_model', '0', model_name, data
+    )
 
     assert output == 'Model retrained successfully'
 
@@ -438,7 +385,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
     output = mlflow_repository.update_production_model_by_run_id('0', 'test')
 
     mlflow.register_model.assert_called_once_with(
-        "runs:/0/prediction_model",
+        'runs:/0/prediction_model',
         'test',
     )
 
@@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
 @patch('model_manager.utils.repository.model_repository.mlflow')
 def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
     mlflow.tracking.MlflowClient.return_value = MagicMock(
-        get_registered_model=MagicMock(
-            return_value=MagicMock(
-                latest_versions={}
-            )
-        )
+        get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
     )
 
     try:
         mlflow_repository.update_production_model_by_run_id('0', 'test')
-    except Exception as e:
+    except Exception as e:  # noqa: BLE001
         assert str(e) == 'Model versions is not a list'
     else:
-        assert False
+        raise AssertionError('Expected exception')
 
 
 def test_update_production_model(mlflow_repository):
     connector = mlflow_repository
 
-    with patch.object(connector, 'get_experiment',
-                      return_value='0') as get_experiment:
-        with patch.object(connector, 'get_experiment_last_run',
-                          return_value='2') as get_experiment_last_run:
-            with patch.object(connector, 'update_production_model_by_run_id',
-                              return_value={'model_name': 'test', 'version': '3',
-                                            'mlflow_run_id': '0'}) as update_production_model_by_run_id:
-
+    with patch.object(connector, 'get_experiment', return_value='0') as get_experiment:
+        with patch.object(
+            connector, 'get_experiment_last_run', return_value='2'
+        ) as get_experiment_last_run:
+            with patch.object(
+                connector,
+                'update_production_model_by_run_id',
+                return_value={'model_name': 'test', 'version': '3', 'mlflow_run_id': '0'},
+            ) as update_production_model_by_run_id:
                 output = connector.update_production_model('0', 'test')
 
                 get_experiment.assert_called_once_with('0')
                 get_experiment_last_run.assert_called_once_with('0')
-                update_production_model_by_run_id.assert_called_once_with(
-                    '2', 'test')
+                update_production_model_by_run_id.assert_called_once_with('2', 'test')
 
                 assert output == {
                     'model_name': 'test',

tests/laborious/utils/test_connectors_config.py

@@ -1,7 +1,10 @@
 from os import environ
-from model_manager.utils.connectors_config import (build_mlflow_config,
-                                               build_postgres_config,
-                                               build_mongodb_config)
+
+from model_manager.utils.connectors_config import (
+    build_mlflow_config,
+    build_mongodb_config,
+    build_postgres_config,
+)
 
 
 def test_build_mlflow_config_with_env_vars():
@@ -95,7 +98,7 @@ def test_build_mongo_db_config_with_env_vars():
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
         'database_name': 'test_db',
-        'ttl_index_seconds': 3600
+        'ttl_index_seconds': 3600,
     }
 
 
@@ -108,5 +111,5 @@ def test_build_mongo_db_config_with_defaults():
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
         'database_name': 'sientia',
-        'ttl_index_seconds': 3600
+        'ttl_index_seconds': 3600,
     }

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -1,9 +1,12 @@
-from unittest.mock import call, patch, AsyncMock, ANY
-from pytest import mark, fixture
+from unittest.mock import ANY, AsyncMock, call, patch
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from model_manager.activities.activities import Activities
-from model_manager.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
+from model_manager.workflows.sub_workflows.format_and_export_prediction import (
+    FormatAndExportPrediction,
+)
 
 
 @fixture
@@ -12,119 +15,133 @@ def format_and_export_prediction():
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
+@patch(
+    'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
 async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
-
     input_data = {
         'metadata': metadata,
-        "path_flag": None,
-        "data": {"test": "data"},
-        "timestamp": "2021-01-01",
-        "model_id": 1,
-        "prediction_confidence": 0,
-        "schema": "test_schema",
-        "table_name": "test_table",
-        "prediction_store_policy": "erl:1"
+        'path_flag': None,
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'prediction_store_policy': 'erl:1',
     }
 
     await format_and_export_prediction.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.format_prediction,
-            {
-                'data': input_data['data'],
-                'timestamp': input_data['timestamp'],
-                'model_id': input_data['model_id'],
-                'prediction_confidence': input_data['prediction_confidence'],
-                'prediction_store_policy': input_data['prediction_store_policy'],
-                **metadata
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_prediction,
+                {
+                    'data': input_data['data'],
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'prediction_store_policy': input_data['prediction_store_policy'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_activity_method.return_value,
-                **metadata,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': workflow_mock.execute_activity_method.return_value,
+                    **metadata,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
     assert workflow_mock.execute_activity_method.call_count == 2
     assert workflow_mock.execute_local_activity_method.call_count == 1
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
+@patch(
+    'model_manager.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
 async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
-
     input_data = {
         'metadata': metadata,
-        "path_flag": "default",
-        "data": {"test": "data"},
-        "timestamp": "2021-01-01",
-        "model_id": 1,
-        "prediction_confidence": 0,
-        "schema": "test_schema",
-        "table_name": "test_table",
-        "comment": "test_comment"
+        'path_flag': 'default',
+        'data': {'test': 'data'},
+        'timestamp': '2021-01-01',
+        'model_id': 1,
+        'prediction_confidence': 0,
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'comment': 'test_comment',
     }
 
     await format_and_export_prediction.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.format_default_prediction,
-            {
-                'timestamp': input_data['timestamp'],
-                'model_id': input_data['model_id'],
-                'prediction_confidence': input_data['prediction_confidence'],
-                'comment': input_data['comment'],
-                **metadata
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_default_prediction,
+                {
+                    'timestamp': input_data['timestamp'],
+                    'model_id': input_data['model_id'],
+                    'prediction_confidence': input_data['prediction_confidence'],
+                    'comment': input_data['comment'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
-            {
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-                'data': workflow_mock.execute_activity_method.return_value,
-                **metadata,
-                'timestamp_conversion': {
-                    'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
-                }
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                    'data': workflow_mock.execute_activity_method.return_value,
+                    **metadata,
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, patch, call, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from model_manager.activities.activities import Activities
 from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
 
@@ -10,17 +12,17 @@ def prediction_process():
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model",
-        "model_name": "test_model",
-        "workflow_name": "test_workflow",
-        "schema_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.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
@@ -34,26 +36,23 @@ async def test_run(workflow_mock, prediction_process):
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
-            'retention': '30'
-        },
+        'model_config': {'retention': '30'},
         'path_priority': ['continue', 'repeat', 'stop'],
-
-        'prediction_store_policy': 'lts:1'
+        'prediction_store_policy': 'lts:1',
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
         {'content': 'predicted_data', 'timestamp': '2024-01-01'},  # request_predict
         # mlflow_response_gate (predict)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
     ]
 
     # Act
@@ -62,57 +61,112 @@ async def test_run(workflow_mock, prediction_process):
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 7
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {
-            **metadata,
-            'data': input_data['data'],
-        },
-            retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.input_gate, {
-            **metadata,
-            'filters': input_data['input_filters'],
-            'data': input_data['data'],
-            'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_transform, {
-            **metadata,
-            'data': input_data['data'],
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            **metadata,
-            'filters': input_data['mlflow_transform_filters'],
-            'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_content_gate, {
-            **metadata,
-            'filters': input_data['mlflow_transform_filters'],
-            'data': 'transformed_data',
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_predict, {
-            **metadata,
-            'data': 'transformed_data',
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            **metadata,
-            'filters': input_data['mlflow_predict_filters'],
-            'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
-            'type': 'predict',
-            'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    **metadata,
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    **metadata,
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': 'transformed_data',
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
+                    **metadata,
+                    'data': 'transformed_data',
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    **metadata,
+                    'filters': input_data['mlflow_predict_filters'],
+                    'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+                    'type': 'predict',
+                    'path_priority': input_data['path_priority'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
     workflow_mock.execute_child_workflow.assert_called_once_with(
         'format_and_export_prediction',
@@ -125,17 +179,16 @@ async def test_run(workflow_mock, prediction_process):
             'model_id': 1,
             'model_name': 'test_model_name',
             'model_config': input_data['model_config'],
-
             'schema': input_data['schema'],
             'table_name': input_data['table_name'],
             'comment': 'Error',
-            'prediction_store_policy': input_data['prediction_store_policy']
-        }
+            'prediction_store_policy': input_data['prediction_store_policy'],
+        },
     )
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.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
@@ -149,17 +202,14 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
-            'retention': '30'
-        },
+        'model_config': {'retention': '30'},
         'path_priority': ['continue', 'repeat', 'stop'],
-
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('stop', 0.95, "Input data with bad quality"),  # input_gate
+        ('stop', 0.95, 'Input data with bad quality'),  # input_gate
     ]
 
     # Act
@@ -167,23 +217,35 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 2
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {
-             'data': input_data['data'],
-             **metadata,
-             }, retry_policy=ANY, start_to_close_timeout=ANY),
-        call(Activities.input_gate, {
-             'filters': input_data['input_filters'],
-             'data': input_data['data'],
-             'path_priority': input_data['path_priority'],
-             **metadata,
-             }, retry_policy=ANY, start_to_close_timeout=ANY)
-    ])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
+                    'data': input_data['data'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.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
@@ -197,19 +259,16 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
-            'retention': '30'
-        },
+        'model_config': {'retention': '30'},
         'path_priority': ['continue', 'repeat', 'stop'],
-
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('repeat', 0.95, "Input data with bad quality"),  # input_gate
+        ('repeat', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (transform)
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (transform)
     ]
 
     # Act
@@ -217,46 +276,72 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 4
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {
-             'data': input_data['data'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.input_gate, {
-             'filters': input_data['input_filters'],
-             'data': input_data['data'],
-             'path_priority': input_data['path_priority'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_transform, {
-            'data': input_data['data'],
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-            **metadata
-        },
-            retry_policy=ANY, start_to_close_timeout=ANY)
-    ])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            'filters': input_data['mlflow_transform_filters'],
-            'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-            **metadata
-        }, retry_policy=ANY, start_to_close_timeout=ANY)
-    ])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
+                    'data': input_data['data'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.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])
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
     # Arrange
     input_data = {
         'metadata': metadata,
@@ -268,22 +353,19 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
-            'retention': '30'
-        },
+        'model_config': {'retention': '30'},
         'path_priority': ['continue', 'repeat', 'stop'],
-
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
     ]
 
     # Act
@@ -292,51 +374,88 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 5
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {
-             'data': input_data['data'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.input_gate, {
-             'filters': input_data['input_filters'],
-             'data': input_data['data'],
-             'path_priority': input_data['path_priority'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_transform, {
-            'data': input_data['data'],
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-            **metadata
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            'filters': input_data['mlflow_transform_filters'],
-            'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-            **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_content_gate, {
-            'filters': input_data['mlflow_transform_filters'],
-            'data': 'transformed_data',
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-            **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
+                    'data': input_data['data'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': 'transformed_data',
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.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])
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
     # Arrange
     input_data = {
         'metadata': metadata,
@@ -348,24 +467,21 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
-            'retention': '30'
-        },
+        'model_config': {'retention': '30'},
         'path_priority': ['continue', 'repeat', 'stop'],
-
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
         {'content': 'predicted_data', 'timestamp': '2024-01-01'},  # request_predict
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (predict)
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (predict)
     ]
 
     # Act
@@ -373,63 +489,117 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 7
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.get_last_timestamp, {
-             'data': input_data['data'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.input_gate, {
-             'filters': input_data['input_filters'],
-             'data': input_data['data'],
-             'path_priority': input_data['path_priority'],
-             **metadata,
-             },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_transform, {
-            'data': input_data['data'],
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-            **metadata
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            'filters': input_data['mlflow_transform_filters'],
-            'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-            **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_content_gate, {
-            'filters': input_data['mlflow_transform_filters'],
-            'data': 'transformed_data',
-            'type': 'transform',
-            'path_priority': input_data['path_priority'],
-            **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.request_predict, {
-            'data': 'transformed_data',
-            'model_name': input_data['model_name'],
-            'model_config': input_data['model_config'],
-            **metadata
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(Activities.mlflow_response_gate, {
-            'filters': input_data['mlflow_predict_filters'],
-            'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
-            'type': 'predict',
-            'path_priority': input_data['path_priority'],
-            **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
+                    'data': input_data['data'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
+                    'filters': input_data['input_filters'],
+                    'data': input_data['data'],
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
+                    'data': input_data['data'],
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
+                    'filters': input_data['mlflow_transform_filters'],
+                    'data': 'transformed_data',
+                    'type': 'transform',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
+                    'data': 'transformed_data',
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
+                    'filters': input_data['mlflow_predict_filters'],
+                    'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
+                    'type': 'predict',
+                    'path_priority': input_data['path_priority'],
+                    **metadata,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_stop(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -440,21 +610,24 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
-        'retention': '30'
-    }
+    model_config = {'retention': '30'}
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
-            'model_config': model_config
-        }, confidence, last_timestamp, ""
+            'model_config': model_config,
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert
@@ -464,7 +637,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -475,21 +648,24 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
-        'retention': '30'
-    }
+    model_config = {'retention': '30'}
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
-            'model_config': model_config
-        }, confidence, last_timestamp, ""
+            'model_config': model_config,
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert
@@ -504,13 +680,13 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
             'last_timestamp': last_timestamp,
         },
         retry_policy=ANY,
-        start_to_close_timeout=ANY
+        start_to_close_timeout=ANY,
     )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -521,14 +697,14 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
-        'retention': '30'
-    }
+    model_config = {'retention': '30'}
     prediction_store_policy = 'erl:1'
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
@@ -536,9 +712,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'last_timestamp': last_timestamp,
             'model_name': model_name,
             'model_config': model_config,
-    
-            'prediction_store_policy': prediction_store_policy
-        }, confidence, last_timestamp, 'Prediction Process'
+            'prediction_store_policy': prediction_store_policy,
+        },
+        confidence,
+        last_timestamp,
+        'Prediction Process',
     )
 
     # Assert
@@ -558,14 +736,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'schema': schema,
             'table_name': table_name,
             'comment': 'Prediction Process',
-    
-            'prediction_store_policy': prediction_store_policy
-        }
+            'prediction_store_policy': prediction_store_policy,
+        },
     )
 
 
 @mark.asyncio
-@patch("model_manager.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('model_manager.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -576,13 +753,13 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
-        'retention': '30'
-    }
+    model_config = {'retention': '30'}
     prediction_store_policy = 'erl:1'
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             **metadata,
             'schema': schema,
             'table_name': table_name,
@@ -590,9 +767,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
             'last_timestamp': last_timestamp,
             'model_name': model_name,
             'model_config': model_config,
-    
-            'prediction_store_policy': prediction_store_policy
-        }, confidence, last_timestamp, ""
+            'prediction_store_policy': prediction_store_policy,
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert

tests/laborious/workflows/test_minimal_retrain.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from model_manager.activities.activities import Activities
 from model_manager.workflows.minimal_retrain import MinimalRetrain
 
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "minimal_retrain",
-        "schedule_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
     },
 }
 
@@ -23,19 +25,19 @@ metadata = {
 @patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
     input_data = {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "minimal_retrain",
-        "schedule_name": "test_schedule",
-        "query": "test_query",
-        "schema": "test_schema",
-        "table_name": "test_table",
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
     }
 
     workflow_mock.execute_activity_method = AsyncMock(
         return_value={
-            "data1": "1",
-            "data2": "2",
+            'data1': '1',
+            'data2': '2',
         }
     )
 
@@ -47,52 +49,58 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
                 Activities.load_custom_query,
                 {
                     **metadata,
-                    "query": input_data["query"],
-                    'datetime_columns': input_data.get('datetime_columns', [])
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
                 },
                 retry_policy=ANY,
-                start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
             )
         ]
     )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.retrain_model,
-            {
-                **metadata,
-                'data': workflow_mock.execute_local_activity_method.return_value,
-                'model_name': input_data['model_name'],
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'model_name': input_data['model_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.update_production_model,
-            {
-                **metadata,
-                'model_name': input_data['model_name'],
-                'model_id': input_data['model_id'],
-                **workflow_mock.execute_activity_method.return_value,
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.update_production_model,
+                {
+                    **metadata,
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    **workflow_mock.execute_activity_method.return_value,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
-        call(
-            Activities.export_data_to_postgres,
-            {
-                **metadata,
-                'data': workflow_mock.execute_activity_method.return_value,
-                'schema': input_data['schema'],
-                'table_name': input_data['table_name'],
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )

tests/laborious/workflows/test_predictions_batch.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, call, patch, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from model_manager.activities.activities import Activities
 from model_manager.workflows.predictions_batch import PredictionsBatch
 
@@ -10,11 +12,11 @@ def predictions_batch() -> PredictionsBatch:
 
 
 metadata = {
-    "metadata": {
-        "model_id": "test_model_id",
-        "model_name": "test_model",
-        "workflow_name": "predictions_batch",
-        "schedule_name": "test_schedule",
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'predictions_batch',
+        'schedule_name': 'test_schedule',
     },
 }
 
@@ -22,9 +24,7 @@ metadata = {
 @mark.asyncio
 @patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
-    workflow_mock.execute_local_activity_method.return_value = {
-        'data': 'test_data'
-    }
+    workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
     input_data = {
         'schedule_name': 'test_schedule',
         'model_name': 'test_model',
@@ -32,28 +32,27 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'query': 'SELECT * FROM test',
         'schema': 'test_schema',
         'table_name': 'test_table',
-
         'datetime_columns': ['timestamp', 'created_at'],
         'prediction_store_policy': 'erl:1',
-        'model_config': {
-            'retention': '30'
-        }
+        'model_config': {'retention': '30'},
     }
 
     await predictions_batch.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
-        call(
-            Activities.load_custom_query,
-            {
-                **metadata,
-                'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
-            },
-            retry_policy=ANY,
-            start_to_close_timeout=ANY
-        )
-    ])
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.load_custom_query,
+                {
+                    **metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     prediction_input = {
         'metadata': metadata,
         'data': {'data': 'test_data'},
@@ -61,28 +60,18 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         'table_name': input_data['table_name'],
         'model_id': input_data['model_id'],
         'model_name': input_data['model_name'],
-        'input_filters': input_data.get('input_filters', {
-            'EMPTY_DATA': {
-                'POLICY': 'STOP'
-            }
-        }),
-        'mlflow_transform_filters': input_data.get('mlflow_transform_filters', {
-            'API_ERROR': {
-                'POLICY': 'STOP'
-            }
-        }),
-        'mlflow_predict_filters': input_data.get('mlflow_predict_filters', {
-            'API_ERROR': {
-                'POLICY': 'STOP'
-            }
-        }),
+        'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
+        'mlflow_transform_filters': input_data.get(
+            'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
+        ),
+        'mlflow_predict_filters': input_data.get(
+            'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
+        ),
         'model_config': input_data.get('model_config', {}),
         'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
-
-        'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
+        'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
     }
 
-    workflow_mock.execute_child_workflow.assert_has_calls([
-        call(
-            'prediction_process', prediction_input)
-    ])
+    workflow_mock.execute_child_workflow.assert_has_calls(
+        [call('prediction_process', prediction_input)]
+    )