Merge pull request #10 from Aignosi/feature/SIENTIAPDE-1241

Feature/SIENTIAPDE-1241: Corrigir bug ao treinar modelos com a flag de normalizar dados ativa

.env.example

@@ -6,8 +6,7 @@ POSTGRES_DBNAME="sientia"
 POSTGRES_MIN_CONNECTIONS="10"
 POSTGRES_MAX_CONNECTIONS="30"
 
-MLFLOW_HOST="http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
-MLFLOW_PORT="80"
+MLFLOW_URL="http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local:80"
 MLFLOW_USERNAME="aignosi"
 MLFLOW_PASSWORD="mlflow_password"
 
@@ -38,11 +37,8 @@ MINIO_READ_TIMEOUT="60"
 # Workflow Activity Timeouts (in seconds)
 # These timeouts are designed to handle large files (up to 200MB)
 TIMEOUT_VALIDATE_PARAMS="30"           # Parameter validation (fast operation)
-TIMEOUT_DOWNLOAD_FILE="600"            # File download from MinIO (10 min for 200MB @ 1MB/s with 3x buffer)
-TIMEOUT_TRAIN_MODEL="1800"             # Model training (30 min for large datasets)
-TIMEOUT_SAVE_MODEL="300"               # Save model to MLFlow (5 min for artifacts upload)
-TIMEOUT_CLEANUP_DIRECTORY="60"         # Cleanup temporary directory (1 min)
-TIMEOUT_DELETE_FILE="60"               # Delete file from MinIO (1 min)
+TIMEOUT_TRAIN_MODEL="2700"             # Model training (30 min for large datasets)
+TIMEOUT_DELETE_FILE="120"               # Delete file from MinIO (1 min)
 TIMEOUT_UPDATE_DATABASE="30"           # Database update operations (30 sec)
 
 EXTRA_PIP_REQUIREMENTS="git+https://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@github.com/Aignosi/sientia-mlops-library.git"

.gitignore

@@ -145,6 +145,7 @@ celerybeat.pid
 env/
 venv/
 ENV/
+venv_311/
 env.bak/
 venv.bak/
 

Makefile

@@ -1,7 +0,0 @@
-VERSION = 1.0.8
-name = sientia-model-manager
-# ENVIRONMENT = production
-
-docker-hub:
-	@docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) .
-	@docker push aignosi.azurecr.io/$(name):$(VERSION)

README.md

@@ -519,26 +519,6 @@ kubectl exec -n temporal <temporal-admin-tools-pod-name> -- \
    
    The Model Manager requires connections to several external services. For local development, you can use the provided port-forward script to establish connections to services running in your Kubernetes cluster.
 
-#### Port Forward Setup Script
-
-The `setup_port_forwards.sh` script automates the creation of port forwards to all required services:
-
-**Features:**
-- 🔄 **Automatic Cleanup**: Kills existing port-forward jobs for the same services
-- ✅ **Port Validation**: Checks if ports are available before creating forwards
-- 🛡️ **Safe Execution**: Stops if any port is already in use by another process
-- 📊 **Clear Output**: Color-coded status messages and service information
-
-**Usage:**
-
-```bash
-# Make script executable (first time only)
-chmod +x setup_port_forwards.sh
-
-# Run the script
-./setup_port_forwards.sh
-```
-
 **Services and Ports:**
 
 | Local Port | Service | Description | Namespace |
@@ -559,17 +539,8 @@ jobs -l
 # Stop all port forwards
 jobs -p | xargs kill
 
-# Stop and restart (using the script)
-./setup_port_forwards.sh
 ```
 
-**Troubleshooting:**
-
-If you encounter port conflicts:
-1. The script will show which ports are in use
-2. Stop the conflicting process or use the script to kill existing port-forwards
-3. Run the script again
-
 **Manual Port Forwarding:**
 
 If you prefer manual control or need different ports:

model_manager/activities/activities.py

@@ -7,12 +7,12 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.observability.logger import Logger
 
     from model_manager.activities.experiment_tracking import ExperimentTracking
-    from model_manager.activities.minio import MinIO
-    from model_manager.activities.mlflow import MLFlow
     from model_manager.activities.training import Training
+    from model_manager.utils.repository.model_repository import ModelRepository
+    from model_manager.utils.repository.storage_repository import StorageRepository
 
 
-class Activities(ExperimentTracking, MLFlow, MinIO, Training):
+class Activities(ExperimentTracking, Training):
     """
     Main activities orchestrator for the Model Manager system.
 
@@ -75,18 +75,14 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Training):
             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'],
+        self.model_repository = ModelRepository(
+            url=mlflow_config['url'],
+            username=mlflow_config['username'],
+            password=mlflow_config['password'],
             logger=logger,
-            notification_handler=notification_handler,
         )
 
-        MinIO.__init__(
-            self,
+        self.storage_repository = StorageRepository(
             endpoint_url=minio_config['endpoint_url'],
             access_key=minio_config['access_key'],
             secret_key=minio_config['secret_key'],
@@ -97,10 +93,15 @@ class Activities(ExperimentTracking, MLFlow, MinIO, Training):
             connect_timeout=minio_config['connect_timeout'],
             read_timeout=minio_config['read_timeout'],
             logger=logger,
-            notification_handler=notification_handler,
         )
 
-        Training.__init__(self, logger=logger, notification_handler=notification_handler)
+        Training.__init__(
+            self,
+            model_repository=self.model_repository,
+            storage_repository=self.storage_repository,
+            logger=logger,
+            notification_handler=notification_handler,
+        )
 
     def __del__(self):
         """

model_manager/activities/experiment_tracking.py

@@ -9,7 +9,9 @@ methods for experiment management.
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
+    import asyncio
     import traceback
+    from collections.abc import Mapping
     from datetime import UTC, datetime
     from enum import Enum
     from typing import Any
@@ -18,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.notifications.models import NotificationLevel
     from sientia_do.observability.logger import Logger
     from sientia_do.temporal.activities.postgres import Postgres
+    from sqlalchemy import text
 
 
 class UpdateType(str, Enum):
@@ -85,8 +88,6 @@ class ExperimentTracking(Postgres):
             logger=logger,
             notification_handler=notification_handler,
         )
-        self.logger = logger
-        self.notification_handler = notification_handler
 
     def __del__(self):
         """
@@ -105,6 +106,118 @@ class ExperimentTracking(Postgres):
                 # Silently ignore errors during garbage collection
                 pass
 
+    async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
+        """
+        Execute an UPDATE SQL statement asynchronously.
+
+        Args:
+            query: Parameterized SQL string to execute.
+            params: Mapping of parameters for the SQL query.
+
+        Returns:
+            dict: A dictionary containing the affected row count: {'rowcount': int}.
+        """
+
+        def _run() -> dict[str, Any]:
+            with self.engine.begin() as connection:
+                result = connection.execute(text(query), params)
+                return {'rowcount': result.rowcount}
+
+        return await asyncio.to_thread(_run)
+
+    def _build_status_update_query(
+        self, status: str | None, experiment_run_id: int
+    ) -> tuple[str, dict[str, Any]]:
+        """Build SQL query for simple status update."""
+        if not isinstance(status, str) or not status:
+            raise ValueError('status is required for STATUS update type')
+
+        sql_query = """
+            UPDATE experiment_run
+            SET status = :status, updated_at = :updated_at
+            WHERE id = :experiment_run_id
+        """
+
+        query_params = {
+            'status': status,
+            'updated_at': datetime.now(UTC),
+            'experiment_run_id': experiment_run_id,
+        }
+
+        return sql_query, query_params
+
+    def _build_status_with_error_query(
+        self, status: str | None, error_message: str | None, experiment_run_id: int
+    ) -> tuple[str, dict[str, Any]]:
+        """Build SQL query for status update with error message."""
+        if not isinstance(status, str) or not status:
+            raise ValueError('status is required for STATUS_WITH_ERROR update type')
+
+        if not isinstance(error_message, str) or not error_message:
+            raise ValueError('error_message is required for STATUS_WITH_ERROR update type')
+
+        # Truncate error message if too long
+        truncated_error = error_message[:1024] if len(error_message) > 1024 else error_message
+
+        sql_query = """
+            UPDATE experiment_run
+            SET status = :status, error_message = :error_message, updated_at = :updated_at
+            WHERE id = :experiment_run_id
+        """
+
+        query_params = {
+            'status': status,
+            'error_message': truncated_error,
+            'updated_at': datetime.now(UTC),
+            'experiment_run_id': experiment_run_id,
+        }
+
+        return sql_query, query_params
+
+    def _build_model_saved_query(
+        self, run_name: str | None, status: str | None, experiment_run_id: int
+    ) -> tuple[str, dict[str, Any]]:
+        """Build SQL query for model saved update."""
+        if not isinstance(run_name, str) or not run_name:
+            raise ValueError('run_name is required for MODEL_SAVED update type')
+
+        if not isinstance(status, str) or not status:
+            raise ValueError('status is required for MODEL_SAVED update type')
+
+        sql_query = """
+            UPDATE experiment_run
+            SET run_name = :run_name, status = :status, updated_at = :updated_at
+            WHERE id = :experiment_run_id
+        """
+
+        query_params = {
+            'run_name': run_name,
+            'status': status,
+            'updated_at': datetime.now(UTC),
+            'experiment_run_id': experiment_run_id,
+        }
+
+        return sql_query, query_params
+
+    def _get_update_query_and_params(
+        self, update_type: str, experiment_run_id: int, input_data: dict[str, Any]
+    ) -> tuple[str, dict[str, Any]]:
+        """Get SQL query and parameters based on update type."""
+        status = input_data.get('status')
+        error_message = input_data.get('error_message')
+        run_name = input_data.get('run_name')
+
+        if update_type == UpdateType.STATUS:
+            return self._build_status_update_query(status, experiment_run_id)
+
+        if update_type == UpdateType.STATUS_WITH_ERROR:
+            return self._build_status_with_error_query(status, error_message, experiment_run_id)
+
+        if update_type == UpdateType.MODEL_SAVED:
+            return self._build_model_saved_query(run_name, status, experiment_run_id)
+
+        raise ValueError(f'Invalid update_type: {update_type}')
+
     @activity.defn(name='update_experiment_run')
     async def update_experiment_run(self, input_data: dict[str, Any]) -> None:
         """
@@ -128,107 +241,33 @@ class ExperimentTracking(Postgres):
         Raises:
             ValueError: If required parameters are missing for the update type
             RuntimeError: If update operation fails
-
-        Example:
-            # Update status only
-            await update_experiment_run({
-                'metadata': {},
-                'experiment_run_id': 123,
-                'update_type': 'status',
-                'status': 'TRAINING_SUCCESS'
-            })
-
-            # Update with error
-            await update_experiment_run({
-                'metadata': {},
-                'experiment_run_id': 123,
-                'update_type': 'status_with_error',
-                'status': 'TRAINING_ERROR',
-                'error_message': 'Model training failed: insufficient data'
-            })
-
-            # Update with model saved
-            await update_experiment_run({
-                'metadata': {},
-                'experiment_run_id': 123,
-                'update_type': 'model_saved',
-                'run_name': 'experiment-model-5'
-            })
         """
         metadata = input_data.get('metadata', {})
         experiment_run_id = input_data['experiment_run_id']
         update_type = input_data['update_type']
         status = input_data.get('status')
-        error_message = input_data.get('error_message')
-        run_name = input_data.get('run_name')
 
         try:
-            self.info(
-                f'Updating experiment run {experiment_run_id} with type: {update_type}', metadata
+            sql_query, query_params = self._get_update_query_and_params(
+                update_type, experiment_run_id, input_data
             )
 
-            # Validate parameters based on update type
-            query_params: tuple[Any, ...]
-            if update_type == UpdateType.STATUS:
-                if not status:
-                    raise ValueError('status is required for STATUS update type')
-                sql_query = """
-                    UPDATE experiment_run
-                    SET status = %s, updated_at = %s
-                    WHERE id = %s
-                """
-                query_params = (status, datetime.now(UTC), experiment_run_id)
-
-            elif update_type == UpdateType.STATUS_WITH_ERROR:
-                if not status or not error_message:
-                    raise ValueError(
-                        'status and error_message are required for STATUS_WITH_ERROR update type'
-                    )
-
-                # Truncate the error_message to 1024 characters if necessary
-                if len(error_message) > 1024:
-                    error_message = error_message[:1024]
-
-                sql_query = """
-                    UPDATE experiment_run
-                    SET status = %s, error_message = %s, updated_at = %s
-                    WHERE id = %s
-                """
-                query_params = (
-                    status,
-                    error_message,
-                    datetime.now(UTC),
-                    experiment_run_id,
-                )
-
-            elif update_type == UpdateType.MODEL_SAVED:
-                if not run_name:
-                    raise ValueError('run_name is required for MODEL_SAVED update type')
-                sql_query = """
-                    UPDATE experiment_run
-                    SET run_name = %s, status = %s, updated_at = %s
-                    WHERE id = %s
-                """
-                query_params = (run_name, status, datetime.now(UTC), experiment_run_id)
-
-            else:
-                raise ValueError(f'Invalid update_type: {update_type}')
-
-            # Execute update query
-            result = await self.execute_query(sql_query, query_params)
+            result = await self._execute_update(sql_query, query_params)
 
             if result.get('rowcount', 0) == 0:
-                error_msg = f'No rows updated for experiment run {experiment_run_id}'
+                error_msg = (
+                    f'No rows updated for experiment run {experiment_run_id} with status {status}'
+                )
                 raise ValueError(error_msg)
 
             self.info(
-                f'Successfully updated experiment run {experiment_run_id} with type {update_type}',
+                f'Successfully updated experiment run {experiment_run_id} with status {status}',
                 metadata,
             )
-
         except Exception as e:  # noqa: BLE001
-            error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Type: {update_type}, Error: {str(e)}'
+            error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Status: {status}, Error: {str(e)}'
             trace = traceback.format_exc()
+
             self.send_notification(
                 metadata=metadata,
                 notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
@@ -237,5 +276,4 @@ class ExperimentTracking(Postgres):
                 level=NotificationLevel.ERROR,
                 attachment_content=trace,
             )
-            self.error(trace, metadata=metadata)
             raise RuntimeError(error_msg) from e

model_manager/activities/minio.py

@@ -1,229 +0,0 @@
-from temporalio import activity, workflow
-
-with workflow.unsafe.imports_passed_through():
-    import traceback
-    from io import BytesIO
-    from typing import Any
-
-    import boto3  # type: ignore[import-untyped]
-    from botocore.config import Config  # type: ignore[import-untyped]
-    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.temporal.activities.base import BaseActivity
-
-
-class MinIO(BaseActivity):
-    """
-    MinIO (S3-compatible) storage activities for file operations.
-
-    This class provides activities for interacting with MinIO object storage,
-    including file download and deletion operations. It handles authentication,
-    connection management, and comprehensive error handling.
-
-    The class implements best practices for S3/MinIO operations:
-    - Connection reuse (boto3 client is thread-safe)
-    - Automatic retry with exponential backoff
-    - Comprehensive error handling and logging
-    - Notification integration for critical errors
-
-    Attributes:
-        endpoint_url (str): MinIO server endpoint URL
-        access_key (str): MinIO access key ID
-        secret_key (str): MinIO secret access key
-        region (str): MinIO region name
-        use_ssl (bool): Whether to use SSL/TLS for connections
-        minio_client: Boto3 S3 client configured for MinIO
-    """
-
-    def __init__(
-        self,
-        endpoint_url: str,
-        access_key: str,
-        secret_key: str,
-        region: str,
-        use_ssl: bool,
-        max_retry_attempts: int,
-        retry_mode: str,
-        connect_timeout: int,
-        read_timeout: int,
-        logger: Logger,
-        notification_handler: NotificationHandler,
-    ):
-        """
-        Initialize MinIO activities with server configuration.
-
-        This constructor creates a persistent boto3 S3 client that will be
-        reused across all activity calls. The client is thread-safe and
-        includes automatic retry configuration.
-
-        Args:
-            endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000)
-            access_key: MinIO access key ID for authentication
-            secret_key: MinIO secret access key for authentication
-            region: MinIO region name (e.g., us-east-1)
-            use_ssl: Whether to use SSL/TLS for connections
-            max_retry_attempts: Maximum number of retry attempts (e.g., 3)
-            retry_mode: Retry mode - standard, legacy, or adaptive (e.g., adaptive)
-            connect_timeout: Connection timeout in seconds (e.g., 10)
-            read_timeout: Read timeout in seconds (e.g., 60)
-            logger: Logger instance for observability and debugging
-            notification_handler: Notification handler for alerts and monitoring
-
-        Raises:
-            ConnectionError: If boto3 client initialization fails
-        """
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
-        self.endpoint_url = endpoint_url
-        self.access_key = access_key
-        self.secret_key = secret_key
-        self.region = region
-        self.use_ssl = use_ssl
-        self.max_retry_attempts = max_retry_attempts
-        self.retry_mode = retry_mode
-        self.connect_timeout = connect_timeout
-        self.read_timeout = read_timeout
-
-        # Configure boto3 with retry strategy
-        # This handles transient network errors and connection issues automatically
-        boto_config = Config(
-            region_name=region,
-            retries={
-                'max_attempts': max_retry_attempts,
-                'mode': retry_mode,
-            },
-            connect_timeout=connect_timeout,
-            read_timeout=read_timeout,
-        )
-
-        try:
-            self.minio_client = boto3.client(
-                's3',
-                endpoint_url=endpoint_url,
-                aws_access_key_id=access_key,
-                aws_secret_access_key=secret_key,
-                config=boto_config,
-                use_ssl=use_ssl,
-            )
-            self.info(f'MinIO client initialized successfully: {endpoint_url}')
-        except Exception as e:
-            error_msg = f'Failed to initialize MinIO client: {str(e)}'
-            self.error(error_msg)
-            raise ConnectionError(error_msg) from e
-
-    @activity.defn(name='fetch_file_from_minio')
-    async def fetch_file_from_minio(self, input_data: dict[str, Any]) -> BytesIO:
-        """
-        Fetch a file from MinIO and return its content as a BytesIO object.
-
-        This activity downloads a file from a MinIO bucket and returns the
-        content as a BytesIO object, which is a file-like object that can be
-        used directly with many Python libraries (pandas, PIL, etc.).
-
-        The operation includes:
-        1. Input validation
-        2. File download from MinIO
-        3. Content reading and wrapping in BytesIO
-        4. Comprehensive error handling and logging
-
-        Args:
-            input_data: Configuration for file fetch operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - bucket_name (str): MinIO bucket name
-                    - file_name (str): File path/key in the bucket
-
-        Returns:
-            BytesIO: File content as a file-like object
-
-        Raises:
-            OSError: If file fetch fails due to network, permission, or other errors
-        """
-        metadata = input_data.get('metadata', {})
-        bucket_name = input_data['bucket_name']
-        file_name = input_data['file_name']
-
-        self.info(f'Fetching file from MinIO: {bucket_name}/{file_name}', metadata)
-
-        try:
-            # Download file from MinIO
-            response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
-
-            # Read file content
-            with response['Body'] as body:
-                file_content = body.read()
-
-            file_size = len(file_content)
-            self.info(
-                f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)',
-                metadata,
-            )
-
-            return BytesIO(file_content)
-
-        except Exception as e:  # noqa: BLE001
-            error_msg = f'Error fetching file from MinIO - Bucket: {bucket_name}, File: {file_name}, Error: {str(e)}'
-            trace = traceback.format_exc()
-            self.send_notification(
-                metadata=metadata,
-                notification_id='FETCH_FILE_FROM_MINIO_ERROR',
-                message=error_msg,
-                block='fetch_file_from_minio',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-            self.error(trace, metadata=metadata)
-            raise OSError(error_msg) from e
-
-    @activity.defn(name='delete_file_from_minio')
-    async def delete_file_from_minio(self, input_data: dict[str, Any]) -> None:
-        """
-        Delete a file from MinIO storage.
-
-        This activity removes a file from a MinIO bucket. The operation is
-        idempotent - deleting a non-existent file is considered successful.
-
-        The operation includes:
-        1. Input validation
-        2. File deletion from MinIO
-        3. Comprehensive error handling and logging
-
-        Args:
-            input_data: Configuration for file deletion operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - bucket_name (str): MinIO bucket name
-                    - file_name (str): File path/key to delete
-
-        Returns:
-            None
-
-        Raises:
-            OSError: If file deletion fails due to permission or other errors
-        """
-        metadata = input_data.get('metadata', {})
-        bucket_name = input_data['bucket_name']
-        file_name = input_data['file_name']
-
-        self.info(f'Deleting file from MinIO: {bucket_name}/{file_name}', metadata)
-
-        try:
-            # Delete file from MinIO
-            # Note: delete_object is idempotent - no error if file doesn't exist
-            self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
-
-            self.info(f'File deleted successfully: {bucket_name}/{file_name}', metadata)
-
-        except Exception as e:  # noqa: BLE001
-            error_msg = f'Error deleting file from MinIO - Bucket: {bucket_name}, File: {file_name}, Error: {str(e)}'
-            trace = traceback.format_exc()
-            self.send_notification(
-                metadata=metadata,
-                notification_id='DELETE_FILE_FROM_MINIO_ERROR',
-                message=error_msg,
-                block='delete_file_from_minio',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-            self.error(trace, metadata=metadata)
-            raise OSError(error_msg) from e

model_manager/activities/mlflow.py

@@ -1,214 +0,0 @@
-from temporalio import activity, workflow
-
-with workflow.unsafe.imports_passed_through():
-    import traceback
-    from typing import Any
-
-    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.temporal.activities.base import BaseActivity
-
-    from model_manager.utils.models.train_model_result import TrainModelResult
-    from model_manager.utils.repository.model_repository import MLFlowRepository
-
-
-class MLFlow(BaseActivity):
-    """
-    MLFlow integration activities for model training operations.
-
-    This class provides activities for saving trained models and managing
-    artifacts in MLFlow. It handles model persistence, artifact generation,
-    and cleanup operations with comprehensive error handling.
-
-    The class implements robust error handling and logging for all
-    MLFlow operations, ensuring reliable model management in production environments.
-
-    Attributes:
-        mlflow_host (str): MLFlow server hostname
-        mlflow_port (int): MLFlow server port
-        mlflow_username (str): MLFlow authentication username
-        mlflow_password (str): MLFlow authentication password
-        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,
-    ):
-        """
-        Initialize MLFlow activities with server configuration.
-
-        Args:
-            mlflow_host: MLFlow server hostname or IP address
-            mlflow_port: MLFlow server port number
-            mlflow_username: Username for MLFlow authentication
-            mlflow_password: Password for MLFlow authentication
-            logger: Logger instance for observability and debugging
-            notification_handler: Notification handler for alerts and monitoring
-
-        Raises:
-            Exception: If MLFlowRepository initialization fails
-        """
-        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
-        )
-
-    @activity.defn(name='save_model')
-    async def save_model(self, input_data: dict[str, Any]) -> TrainModelResult:
-        """
-        Save a trained ML model and its artifacts to MLflow.
-
-        This activity orchestrates the complete model saving pipeline:
-        1. Generates the next run name for the experiment
-        2. Creates and organizes artifacts (reports, data files)
-        3. Logs model, parameters, metrics, and artifacts to MLflow
-
-        Args:
-            input_data: Configuration for model saving operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - train_result (TrainModelResult): Training result with model and metrics
-
-        Returns:
-            TrainModelResult: Updated training result with run_name and artifacts
-
-        Raises:
-            Exception: If model saving fails (after sending notification)
-
-        Example:
-            result = await save_model({
-                'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
-                'train_result': TrainModelResult(...)
-            })
-            # Returns: TrainModelResult with run_name and artifacts
-        """
-        metadata = input_data.get('metadata', {})
-        train_result = input_data['train_result']
-
-        try:
-            experiment_name = train_result.params.experiment_name
-
-            self.info(
-                f'Starting model save for experiment: {experiment_name}',
-                metadata,
-            )
-
-            # Step 1: Generate next run name
-            self.info('Generating run name', metadata)
-            train_result.run_name = self.model_monitoring_repository.get_next_run_name(
-                experiment_name
-            )
-            self.info(f'Generated run name: {train_result.run_name}', metadata)
-
-            # Step 2: Generate artifacts (reports, CSV files)
-            self.info('Generating artifacts', metadata)
-            train_result = self.model_monitoring_repository.generate_artifacts(train_result)
-            self.info('Artifacts generated successfully', metadata)
-
-            # Step 3: Save run to MLflow
-            self.info('Saving run to MLflow', metadata)
-            self.model_monitoring_repository.save_run(train_result)
-
-            self.info(
-                f'Model saved successfully - Run: {train_result.run_name}, '
-                f'Experiment: {experiment_name}',
-                metadata,
-            )
-
-            return train_result
-
-        except Exception as e:  # noqa: BLE001
-            error_msg = f'Error saving model - Experiment: {train_result.params.experiment_name if train_result and train_result.params else "unknown"}, Error: {str(e)}'
-            trace = traceback.format_exc()
-
-            # Send notification (MongoDB)
-            self.send_notification(
-                metadata=metadata,
-                notification_id='SAVE_MODEL_ERROR',
-                message=error_msg,
-                block='save_model',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-
-            # Log error with metadata
-            self.error(trace, metadata=metadata)
-
-            # Re-raise exception to stop workflow
-            raise
-
-    @activity.defn(name='cleanup_run_directory')
-    async def cleanup_run_directory(self, input_data: dict[str, Any]) -> None:
-        """
-        Clean up temporary run directory after model training.
-
-        This activity deletes the temporary directory created during model training
-        and artifact generation. It implements idempotent cleanup to handle cases
-        where the directory may have already been deleted.
-
-        Args:
-            input_data: Configuration for cleanup operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - run_dir (str): Path to the run directory to delete
-
-        Raises:
-            Exception: If cleanup fails for reasons other than directory not existing
-
-        Example:
-            await cleanup_run_directory({
-                'metadata': {'workflow_id': 'cleanup-123'},
-                'run_dir': '/path/to/run_dir'
-            })
-        """
-        import os
-        import shutil
-
-        metadata = input_data.get('metadata', {})
-        run_dir = input_data.get('run_dir')
-
-        try:
-            if not run_dir:
-                self.info('No run directory specified, skipping cleanup', metadata)
-                return
-
-            self.info(f'Cleaning up run directory: {run_dir}', metadata)
-
-            # Idempotent cleanup: check if directory exists before deleting
-            if os.path.exists(run_dir):
-                shutil.rmtree(run_dir)
-                self.info(f'Run directory deleted successfully: {run_dir}', metadata)
-            else:
-                self.info(f'Run directory already deleted: {run_dir}', metadata)
-
-        except Exception as e:
-            error_msg = f'Error cleaning up run directory {run_dir}: {str(e)}'
-            trace = traceback.format_exc()
-
-            # Send notification
-            self.send_notification(
-                metadata=metadata,
-                notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
-                message=error_msg,
-                block='cleanup_run_directory',
-                level=NotificationLevel.ERROR,
-                attachment_content=trace,
-            )
-
-            # Log error
-            self.error(trace, metadata=metadata)
-
-            # Re-raise exception
-            raise

model_manager/activities/training.py

@@ -10,7 +10,6 @@ from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
     import traceback
-    from io import BytesIO
     from typing import Any
 
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
@@ -18,8 +17,10 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.observability.logger import Logger
     from sientia_do.temporal.activities.base import BaseActivity
 
+    from model_manager.utils.exceptions import ModelTrainingError
     from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
+    from model_manager.utils.repository.model_repository import ModelRepository
+    from model_manager.utils.repository.storage_repository import StorageRepository
     from model_manager.utils.repository.training_repository import TrainingRepository
 
 
@@ -39,6 +40,8 @@ class Training(BaseActivity):
 
     def __init__(
         self,
+        model_repository: ModelRepository,
+        storage_repository: StorageRepository,
         logger: Logger,
         notification_handler: NotificationHandler,
     ):
@@ -49,8 +52,10 @@ class Training(BaseActivity):
             logger: Logger instance for observability
             notification_handler: Handler for sending notifications
         """
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
+        super().__init__(logger, notification_handler, set_error_counter=True)
         self.training_repository = TrainingRepository(logger)
+        self.model_repository = model_repository
+        self.storage_repository = storage_repository
 
     @activity.defn(name='validate_train_params')
     async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
@@ -71,27 +76,11 @@ class Training(BaseActivity):
 
         Raises:
             ValueError, TypeError, KeyError: If validation fails (after sending notification)
-
-        Example:
-            result = await validate_train_params({
-                'metadata': {'workflow_id': 'train-123'},
-                'experiment_run_id': 456,
-                'target_variable': 'price',
-                'variable_columns': ['feature1', 'feature2'],
-                'train_size': 80,
-                # ... other required fields at same level
-            })
-            # Returns: TrainModelParams(...)
         """
         metadata = input_data.get('metadata', {})
 
         try:
-            self.info('Validating training parameters', metadata)
-
-            # Step 1: Convert input_data to TrainModelParams (validates types and required fields)
             train_params = TrainModelParams.from_dict(input_data)
-
-            # Step 2: Validate business rules (ranges, consistency, etc.)
             train_params.validate_business_rules()
 
             self.info(
@@ -102,12 +91,10 @@ class Training(BaseActivity):
             )
 
             return train_params
-
         except (ValueError, TypeError, KeyError) as e:
             error_msg = f'Error validating training parameters: {str(e)}'
             trace = traceback.format_exc()
 
-            # Send notification (MongoDB)
             self.send_notification(
                 metadata=metadata,
                 notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
@@ -116,91 +103,68 @@ class Training(BaseActivity):
                 level=NotificationLevel.ERROR,
                 attachment_content=trace,
             )
-
-            # Log error with metadata
-            self.error(trace, metadata=metadata)
-
-            # Re-raise exception to stop workflow
             raise
 
     @activity.defn(name='train_model')
-    async def train_model(self, input_data: dict[str, Any]) -> TrainModelResult:
+    async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
         """
         Train a machine learning model.
 
-        This activity orchestrates the complete ML training pipeline:
-        1. Validates input parameters
-        2. Trains the model using TrainingRepository
-        3. Performs post-training calculations
+        This activity orchestrates the ML training pipeline:
+        1. Validate input parameters.
+        2. Train the model via TrainingRepository.
+        3. Perform post-training calculations.
 
         Args:
-            input_data: Configuration for model training operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
-                    - train_params (TrainModelParams): Training parameters object
+            input_data: Training configuration containing:
+                - metadata (dict): Workflow execution metadata.
+                - uploaded_file (BytesIO): Training data already downloaded from MinIO.
+                - train_params (TrainModelParams | dict): Training parameters.
 
         Returns:
-            TrainModelResult: Training result with model, metrics, and data
+            dict: Keys `run_name` and `run_dir` when training and saving succeed.
 
         Raises:
-            ValueError: If input validation fails
-            Exception: If training fails (after sending notification)
-
-        Example:
-            result = await train_model({
-                'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
-                'uploaded_file': BytesIO(csv_data),
-                'train_params': TrainModelParams(...)
-            })
-            # Returns: TrainModelResult(...)
+            ValueError: If input validation fails.
+            Exception: If training fails (after sending notification).
         """
         metadata = input_data.get('metadata', {})
-        uploaded_file = input_data['uploaded_file']
         train_params = input_data['train_params']
 
+        if isinstance(train_params, dict):
+            train_params = TrainModelParams.from_dict(train_params)
+        # type: ignore[assignment]
+
+        model_trained = False
+        model_saved = False
+
         try:
-            # Validate uploaded_file is BytesIO
-            if not isinstance(uploaded_file, BytesIO):
-                raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
+            with self.storage_repository.fetch_file(
+                train_params.bucket_name, train_params.file_name
+            ) as uploaded_file:
+                train_result = self.training_repository.train(uploaded_file, train_params)
 
-            # Validate train_params is TrainModelParams
-            if not isinstance(train_params, TrainModelParams):
-                raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')
+                train_result = self.training_repository.after_train_calculation(
+                    train_params, train_result
+                )
 
-            self.info(
-                f'Starting model training for target: {train_params.target_variable}',
-                metadata,
-            )
-
-            # Step 1: Train the model
-            self.info('Training model with TrainingRepository', metadata)
-            train_result = self.training_repository.train(uploaded_file, train_params)
-
-            # Step 2: Perform post-training calculations
-            self.info('Performing post-training calculations', metadata)
-            final_result = self.training_repository.after_train_calculation(
-                train_params, train_result
-            )
-
-            self.info(
-                f'Model training completed successfully - '
-                f'MSE: {final_result.mse_val}, MAE: {final_result.mae_val}, R²: {final_result.r2_val}',
-                metadata,
-            )
-
-            return final_result
+                model_trained = True
+                train_result = self.model_repository.save_model(train_result)
+                model_saved = True
 
+                return {
+                    'run_name': train_result.run_name,
+                    'run_dir': train_result.run_dir,
+                }
         except Exception as e:  # noqa: BLE001
-            target = (
-                train_params.target_variable
-                if hasattr(train_params, 'target_variable')
-                else 'unknown'
+            error_msg = (
+                'Error training model - '
+                f'model_trained={model_trained}, model_saved={model_saved}, '
+                f'error: {str(e)}'
             )
-            error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
+
             trace = traceback.format_exc()
 
-            # Send notification (MongoDB)
             self.send_notification(
                 metadata=metadata,
                 notification_id='TRAIN_MODEL_ERROR',
@@ -210,8 +174,48 @@ class Training(BaseActivity):
                 attachment_content=trace,
             )
 
-            # Log error with metadata
-            self.error(trace, metadata=metadata)
+            raise ModelTrainingError(
+                model_trained=model_trained,
+                model_saved=model_saved,
+            ) from e
 
-            # Re-raise exception to stop workflow
+    @activity.defn(name='cleanup_resources')
+    async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
+        """
+        Cleanup temporary resources created during training.
+
+        Args:
+            input_data: Cleanup configuration containing:
+                - metadata (dict): Workflow execution metadata.
+                - run_dir (str): Temporary directory to remove.
+                - bucket_name (str): MinIO bucket of the uploaded file.
+                - file_name (str): MinIO object key to delete.
+
+        Raises:
+            Exception: If cleanup fails (after sending notification).
+        """
+        metadata = input_data.get('metadata', {})
+        run_dir = input_data.get('run_dir', '')
+        bucket_name = input_data.get('bucket_name', '')
+        file_name = input_data.get('file_name', '')
+
+        try:
+            self.model_repository.cleanup_run_directory(run_dir)
+            self.storage_repository.delete_file(bucket_name, file_name)
+        except Exception as e:  # noqa: BLE001
+            error_msg = (
+                f'Error cleaning up resources - Run directory: {run_dir}, '
+                f'File: {bucket_name}/{file_name}, Error: {str(e)}'
+            )
+
+            trace = traceback.format_exc()
+
+            self.send_notification(
+                metadata=metadata,
+                notification_id='CLEANUP_RESOURCES_ERROR',
+                message=error_msg,
+                block='cleanup_resources',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
             raise

model_manager/sientia/model_serving.py

@@ -31,7 +31,6 @@ class ModelServing:
         tracking_uri: str,
         username: str | None = None,
         password: str | None = None,
-        logger: Any | None = None,
     ):
         """
         Initialize ModelServing client.
@@ -54,9 +53,6 @@ class ModelServing:
         if password is not None:
             os.environ['MLFLOW_TRACKING_PASSWORD'] = password
 
-        # Note: logger parameter is accepted but not used
-        # Consider removing if not needed, or implement logging
-
     # Function to list runs for a given experiment
     def search_runs_by_name(
         self, experiment_names: list[str], order_by: None | list[str] = None
@@ -70,8 +66,8 @@ class ModelServing:
         Returns:
             pandas.DataFrame: A DataFrame containing run information.
 
-        Raise:
-            SientiaMlException if unable to search runs
+        Raises:
+            SientiaMlException: If unable to search runs.
         """
         try:
             runs = mlflow.search_runs(experiment_names=experiment_names, order_by=order_by)
@@ -86,6 +82,9 @@ class ModelServing:
 
         Args:
             experiment_identifier (str): name or id of the experiment to be setted
+
+        Raises:
+            Exception: If setting the experiment fails.
         """
         mlflow.set_experiment(experiment_identifier)
 
@@ -103,6 +102,9 @@ class ModelServing:
         Security Warning:
             The GitHub token is hardcoded. Consider moving to environment variable
             or using a secure secret management solution (e.g., K8s secrets).
+
+        Raises:
+            Exception: If logging the model fails.
         """
         mlflow.sklearn.log_model(
             sk_model,
@@ -121,6 +123,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the parameter fails.
         """
         mlflow.log_param(key, value)
 
@@ -134,6 +139,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the metric fails.
         """
         mlflow.log_metric(key, value)
 
@@ -150,6 +158,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the artifact fails.
         """
         mlflow.log_artifact(local_path=local_path, artifact_path=artifact_path, run_id=run_id)
 
@@ -182,11 +193,8 @@ class ModelServing:
         Yields:
             ActiveRun: object that acts as a context manager wrapping the run's state.
 
-        Example:
-            with model_serving.save_experiment(run_name="my_run") as run:
-                model_serving.log_param("param1", value1)
-                model_serving.log_metric("metric1", value2)
-                # Run is automatically closed here, even if an exception occurs
+        Raises:
+            Exception: If starting or ending the MLflow run fails.
         """
         run = mlflow.start_run(
             run_id=run_id,

model_manager/sientia/models.py

@@ -535,8 +535,9 @@ class DataPreprocessor(BaseEstimator, TransformerMixin):
 
             # Normalization
             if step == NORMALIZATION and self.scaler:
-                data_treat = data_treat[self.feature_names_order]
-                data_treat[existing_columns] = self.scaler.transform(data_treat[existing_columns])
+                # Only transform feature columns, preserve target and any other required columns
+                feature_cols = self.feature_names_order
+                data_treat[feature_cols] = self.scaler.transform(data_treat[feature_cols])
 
             # Feature Creation
             if step == 'Feature Creation':

model_manager/utils/connectors_config.py

@@ -42,8 +42,7 @@ def build_mlflow_config() -> dict[str, Any]:
     It handles server connection and authentication parameters.
 
     Environment Variables:
-        MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
-        MLFLOW_PORT: MLFlow server port (default: 5080)
+        MLFLOW_URL: MLFlow server hostname (default: http://localhost:5080)
         MLFLOW_USERNAME: MLFlow username (default: aignosi)
         MLFLOW_PASSWORD: MLFlow password (default: aignosi)
 
@@ -51,8 +50,7 @@ def build_mlflow_config() -> dict[str, Any]:
         dict: MLFlow configuration dictionary with all required parameters
     """
     return {
-        'host': getenv('MLFLOW_HOST', 'http://localhost'),
-        'port': int(getenv('MLFLOW_PORT', '5080')),
+        'url': getenv('MLFLOW_URL', 'http://localhost:5080'),
         'username': getenv('MLFLOW_USERNAME', 'aignosi'),
         'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
     }

model_manager/utils/exceptions.py

@@ -0,0 +1,38 @@
+"""
+Custom exception types for the Model Manager.
+
+This module defines domain-specific exceptions used across the training
+workflow to convey additional context (e.g., flags indicating which steps
+completed successfully) without altering control flow semantics.
+"""
+
+
+class ModelTrainingError(Exception):
+    """
+    Exception raised when the model training workflow fails.
+
+    This exception carries flags indicating whether the model was trained
+    and/or saved successfully, enabling the workflow to map errors to
+    appropriate experiment statuses.
+    """
+
+    def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None):
+        """
+        Initialize ModelTrainingError with training state flags.
+
+        Args:
+            model_trained: True if the training step completed successfully.
+            model_saved: True if the model saving step completed successfully.
+            message: Optional custom error message. If None, a default message
+                including the state flags is generated.
+        """
+        self.model_trained = model_trained
+        self.model_saved = model_saved
+
+        if message is None:
+            message = (
+                'Model training workflow failed '
+                f'(model_trained={model_trained}, model_saved={model_saved})'
+            )
+
+        super().__init__(message)

model_manager/utils/logger_helper.py

@@ -0,0 +1,27 @@
+"""
+Logger helper to prevent duplicate logs caused by propagation.
+
+This module provides a wrapper around sientia_do Logger to disable
+log propagation and prevent duplicate log entries in the Model Manager.
+"""
+
+from sientia_do.observability.logger import Logger as SientiaLogger
+
+
+def get_logger(name: str) -> SientiaLogger:
+    """
+    Create a Logger instance with propagation disabled.
+
+    This prevents duplicate logs caused by hierarchical propagation
+    in Python's logging system.
+
+    Args:
+        name: Logger name (typically __name__ of the calling module).
+
+    Returns:
+        Logger: Configured logger instance with propagation disabled.
+    """
+    logger = SientiaLogger(name)
+    # Disable propagation to prevent duplicate logs
+    logger.base_logger.propagate = False
+    return logger

model_manager/utils/models/experiment_status.py

@@ -14,6 +14,7 @@ class ExperimentStatus(str, Enum):
     to maintain compatibility with existing database records and monitoring systems.
 
     Attributes:
+        ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation.
         MAGE_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
         TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
         TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.

model_manager/utils/models/train_model_params.py

@@ -78,14 +78,6 @@ class TrainModelParams:
             ValueError: If any required field is missing or None
             TypeError: If any field has an incorrect type
             KeyError: If any required key is missing from the dictionary
-
-        Example:
-            >>> input_data = {
-            ...     'variable_columns': ['var1', 'var2'],
-            ...     'lag_train': 5,
-            ...     # ... other fields
-            ... }
-            >>> params = TrainModelParams.from_dict(input_data)
         """
         return cls(
             variable_columns=cls._check_none(
@@ -179,27 +171,23 @@ class TrainModelParams:
 
         Raises:
             ValueError: If any business rule is violated
-
-        Example:
-            >>> params = TrainModelParams.from_dict(data)
-            >>> params.validate_business_rules()  # Raises ValueError if invalid
         """
-        # Validate train_size range (1-99%)
-        if not 1 <= self.train_size <= 99:
-            raise ValueError(f'train_size must be between 1 and 99, got {self.train_size}')
+        # Validate train_size range (10-100%)
+        if not 10 <= self.train_size <= 100:
+            raise ValueError(f'train_size must be between 10 and 100, got {self.train_size}')
 
         # Validate variable_columns is not empty
         if not self.variable_columns:
             raise ValueError('variable_columns cannot be empty')
 
         # Validate positive integers
-        if self.lag_train <= 0:
+        if self.lag_train < 0:
             raise ValueError(f'lag_train must be positive, got {self.lag_train}')
 
-        if self.lag_val <= 0:
+        if self.lag_val < 0:
             raise ValueError(f'lag_val must be positive, got {self.lag_val}')
 
-        if self.window <= 0:
+        if self.window < 0:
             raise ValueError(f'window must be positive, got {self.window}')
 
         # Validate low_lim and upp_lim consistency
@@ -218,13 +206,6 @@ class TrainModelParams:
                     f'Got low_lim={self.low_lim[var]}, upp_lim={self.upp_lim[var]}'
                 )
 
-        # Validate target_variable is in variable_columns
-        if self.target_variable not in self.variable_columns:
-            raise ValueError(
-                f'target_variable "{self.target_variable}" must be in variable_columns: '
-                f'{self.variable_columns}'
-            )
-
         # Validate bucket_name and file_name are not empty
         if not self.bucket_name.strip():
             raise ValueError('bucket_name cannot be empty or whitespace')

model_manager/utils/models/train_model_result.py

@@ -28,6 +28,8 @@ class TrainModelResult:
         mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
         mae_val (float | None): The Mean Absolute Error (MAE) of the predictions. Default is None.
         r2_val (float | None): The R-squared (R²) value of the predictions. Default is None.
+        equation (dict | None): The equation of the model. Default is None.
+        equation_path (str | None): The path to the equation file. Default is None.
         run_name (str | None): The name of the MLFlow run. Default is None.
         report_path (str | None): The path to the generated HTML report file. Default is None.
         train_data_path (str | None): The path to the training dataset CSV file. Default is None.
@@ -47,6 +49,8 @@ class TrainModelResult:
     mse_val: float | None = None
     mae_val: float | None = None
     r2_val: float | None = None
+    equation: dict | None = None
+    equation_path: str | None = None
     run_name: str | None = None
     report_path: str | None = None
     train_data_path: str | None = None

model_manager/utils/repository/model_repository.py

@@ -9,7 +9,10 @@ and logging model runs to MLFlow.
 
 """
 
+import json
+import os
 import shutil
+import warnings
 from datetime import datetime
 from os import makedirs, path
 
@@ -21,15 +24,72 @@ from model_manager.sientia.model_serving import ModelServing  # type: ignore[imp
 from model_manager.sientia.reports import Reports  # type: ignore[import-untyped]
 from model_manager.utils.models.train_model_result import TrainModelResult
 
+# Suppress sklearn FutureWarning about 'squared' deprecation without changing business logic
+warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' is deprecated.*")
+
+
+class ModelRepository:
+    def __init__(self, url, username, password, logger: Logger):
+        self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
 
-class MLFlowRepository:
-    def __init__(self, host, username, password, logger: Logger):
-        self.model_serving = ModelServing(
-            tracking_uri=host, username=username, password=password, logger=logger
-        )
         self.logger = logger
 
-    def get_next_run_name(self, experiment_name: str) -> str:
+    def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
+        """
+        Save a trained ML model and its artifacts to MLflow.
+
+        This activity orchestrates the complete model saving pipeline:
+        1. Generates the next run name for the experiment
+        2. Creates and organizes artifacts (reports, data files)
+        3. Logs model, parameters, metrics, and artifacts to MLflow
+
+        Args:
+            input_data: Configuration for model saving operation
+                Required keys:
+                    - metadata (dict): Workflow execution metadata
+                    - train_result (TrainModelResult): Training result with model and metrics
+
+        Returns:
+            TrainModelResult: Updated training result with run_name and artifacts
+
+        Raises:
+            Exception: If model saving fails (after sending notification)
+        """
+        experiment_name = train_result.params.experiment_name
+        train_result.run_name = self._get_next_run_name(experiment_name)
+        train_result = self._generate_artifacts(train_result)
+        self._save_run(train_result)
+
+        self.logger.info(
+            f'Model saved successfully - experiment run id: {train_result.params.experiment_run_id}, '
+            f'experiment name: {experiment_name}, '
+            f'run name: {train_result.run_name}'
+        )
+
+        return train_result
+
+    def cleanup_run_directory(self, run_dir: str) -> None:
+        """
+        Clean up temporary run directory after model training.
+
+        This activity deletes the temporary directory created during model training
+        and artifact generation. It implements idempotent cleanup to handle cases
+        where the directory may have already been deleted.
+
+        Args:
+            run_dir (str): Path to the run directory to delete
+        """
+        if not run_dir:
+            self.logger.info('No run directory specified, skipping cleanup')
+            return
+
+        if os.path.exists(run_dir):
+            shutil.rmtree(run_dir)
+            self.logger.info(f'Run directory deleted successfully: {run_dir}')
+        else:
+            self.logger.info(f'Run directory already deleted: {run_dir}')
+
+    def _get_next_run_name(self, experiment_name: str) -> str:
         """
         Generates the next run name for a given experiment.
 
@@ -46,7 +106,7 @@ class MLFlowRepository:
         next_run_number = len(runs) + 1
         return f'{experiment_name}-{next_run_number}'
 
-    def generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
+    def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
         """
         Generates and organizes artifacts related to the training process, such as reports and data files.
 
@@ -87,7 +147,7 @@ class MLFlowRepository:
         self._setup_run_directory(data.run_dir, header_file_path)
         return self._generate_report(reference_data, current_data, data)
 
-    def save_run(self, data: TrainModelResult):
+    def _save_run(self, data: TrainModelResult):
         """
         Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
         to the Sientia tracking system.
@@ -122,60 +182,52 @@ class MLFlowRepository:
             self.logger.error(error_msg)
             raise ValueError(error_msg)
 
-        try:
-            # Prepare parameters
-            train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
-            interval_strs = [
-                (str(interval[0]), str(interval[1]))
-                for interval in (data.params.removed_intervals or [])
-            ]
+        # Prepare parameters
+        train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
 
-            # Set experiment and create run
-            self.model_serving.set_experiment(data.params.experiment_name)
-            self.logger.info(
-                f"Logging run '{data.run_name}' to experiment '{data.params.experiment_name}'"
-            )
+        interval_strs = [
+            (str(interval[0]), str(interval[1]))
+            for interval in (data.params.removed_intervals or [])
+        ]
 
-            with self.model_serving.save_experiment(
-                run_name=data.run_name, description=data.params.experiment_name
-            ):
-                # Log model parameters
-                self.model_serving.log_param('model_type', 'Linear Regression')
-                self.model_serving.log_param('target_variable', data.params.target_variable)
-                self.model_serving.log_param('input_variables', data.params.variable_columns)
-                self.model_serving.log_param('lag_train', data.params.lag_train)
-                self.model_serving.log_param('lag_val', data.params.lag_val)
-                self.model_serving.log_param('ma', data.params.window)
-                self.model_serving.log_param('low_lim', data.params.low_lim)
-                self.model_serving.log_param('upp_lim', data.params.upp_lim)
-                self.model_serving.log_param('normalized', data.scaler_dict)
-                self.model_serving.log_param('ar', data.params.include_ar)
-                self.model_serving.log_param('Train_test_split', train_test_split)
-                self.model_serving.log_param('Removed_intervals', interval_strs)
-                self.model_serving.log_param('Retrain', False)
+        # Set experiment and create run
+        self.model_serving.set_experiment(data.params.experiment_name)
 
-                # Log evaluation metrics
-                self.model_serving.log_metric('MSE', data.mse_val)
-                self.model_serving.log_metric('R2', data.r2_val)
-                self.model_serving.log_metric('MAE', data.mae_val)
+        with self.model_serving.save_experiment(
+            run_name=data.run_name, description=data.params.experiment_name
+        ):
+            # Log model parameters
+            self.model_serving.log_param('model_type', 'Linear Regression')
+            self.model_serving.log_param('target_variable', data.params.target_variable)
+            self.model_serving.log_param('input_variables', data.params.variable_columns)
+            self.model_serving.log_param('lag_train', data.params.lag_train)
+            self.model_serving.log_param('lag_val', data.params.lag_val)
+            self.model_serving.log_param('ma', data.params.window)
+            self.model_serving.log_param('low_lim', data.params.low_lim)
+            self.model_serving.log_param('upp_lim', data.params.upp_lim)
+            self.model_serving.log_param('normalized', data.scaler_dict)
+            self.model_serving.log_param('ar', data.params.include_ar)
+            self.model_serving.log_param('Train_test_split', train_test_split)
+            self.model_serving.log_param('Removed_intervals', interval_strs)
+            self.model_serving.log_param('Retrain', False)
 
-                # Log models
-                self.model_serving.log_model(data.process_data, 'data_model')
-                self.model_serving.log_model(data.regr, 'prediction_model')
+            # Log evaluation metrics
+            self.model_serving.log_metric('MSE', data.mse_val)
+            self.model_serving.log_metric('R2', data.r2_val)
+            self.model_serving.log_metric('MAE', data.mae_val)
 
-                # Log artifacts
-                self.model_serving.log_artifact(data.report_path)
-                self.model_serving.log_artifact(data.train_data_path)
-                self.model_serving.log_artifact(data.test_data_path)
+            # Log models
+            self.model_serving.log_model(data.process_data, 'data_model')
+            self.model_serving.log_model(data.regr, 'prediction_model')
 
-                self.logger.info(
-                    f"Successfully logged run '{data.run_name}' with metrics: MSE={data.mse_val:.4f}, R2={data.r2_val:.4f}, MAE={data.mae_val:.4f}"
-                )
+            # Log artifacts
+            self.model_serving.log_artifact(data.report_path)
+            self.model_serving.log_artifact(data.train_data_path)
+            self.model_serving.log_artifact(data.test_data_path)
 
-        except Exception as e:
-            error_msg = f"Failed to save run '{data.run_name}' to MLflow: {str(e)}"
-            self.logger.error(error_msg)
-            raise RuntimeError(error_msg) from e
+            # Log equation artifact if available
+            if data.equation_path and path.exists(data.equation_path):
+                self.model_serving.log_artifact(data.equation_path)
 
     def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
         """
@@ -258,7 +310,6 @@ class MLFlowRepository:
 
         try:
             makedirs(run_dir, exist_ok=True)
-            self.logger.info(f'Created run directory: {run_dir}')
             return run_dir
         except PermissionError as e:
             error_msg = f'Permission denied when creating directory: {run_dir}'
@@ -298,9 +349,6 @@ class MLFlowRepository:
             # Copy header file to run directory
             header_dest = path.join(run_dir, 'header.html')
             shutil.copy(header_file_path, header_dest)
-
-            self.logger.info(f'Run directory setup completed successfully in: {run_dir}')
-
         except FileNotFoundError as e:
             error_msg = f'Header file not found: {header_file_path}'
             self.logger.error(error_msg)
@@ -360,20 +408,22 @@ class MLFlowRepository:
             # Save HTML report
             data.report_path = path.join(data.run_dir, 'report.html')
             report.save_all_sections_html(data.report_path)
-            self.logger.info(f'Generated HTML report: {data.report_path}')
 
             # Save training data CSV
             data.train_data_path = path.join(data.run_dir, 'train_data.csv')
             reference_data.to_csv(data.train_data_path, index=False)
-            self.logger.info(f'Saved training data: {data.train_data_path}')
 
             # Save test data CSV
             data.test_data_path = path.join(data.run_dir, 'test_data.csv')
             current_data.to_csv(data.test_data_path, index=False)
-            self.logger.info(f'Saved test data: {data.test_data_path}')
+
+            # Save equation as JSON
+            if data.equation is not None:
+                data.equation_path = path.join(data.run_dir, 'model_equation.json')
+                with open(data.equation_path, 'w', encoding='utf-8') as f:
+                    json.dump(data.equation, f, indent=2, ensure_ascii=False)
 
             return data
-
         except ValueError as e:
             error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
             self.logger.error(error_msg)

model_manager/utils/repository/storage_repository.py

@@ -0,0 +1,127 @@
+from io import BytesIO
+
+import boto3  # type: ignore[import-untyped]
+from botocore.config import Config  # type: ignore[import-untyped]
+from sientia_do.observability.logger import Logger
+
+
+class StorageRepository:
+    """
+    MinIO (S3-compatible) storage activities for file operations.
+
+    This class provides activities for interacting with MinIO object storage,
+    including file download and deletion operations. It handles authentication,
+    connection management, and comprehensive error handling.
+
+    The class implements best practices for S3/MinIO operations:
+    - Connection reuse (boto3 client is thread-safe)
+    - Automatic retry with exponential backoff
+    - Comprehensive error handling and logging
+    - Notification integration for critical errors
+
+    Attributes:
+        endpoint_url (str): MinIO server endpoint URL
+        access_key (str): MinIO access key ID
+        secret_key (str): MinIO secret access key
+        region (str): MinIO region name
+        use_ssl (bool): Whether to use SSL/TLS for connections
+        minio_client: Boto3 S3 client configured for MinIO
+    """
+
+    def __init__(
+        self,
+        endpoint_url: str,
+        access_key: str,
+        secret_key: str,
+        region: str,
+        use_ssl: bool,
+        max_retry_attempts: int,
+        retry_mode: str,
+        connect_timeout: int,
+        read_timeout: int,
+        logger: Logger,
+    ):
+        """
+        Initialize a reusable MinIO client with retry configuration.
+
+        Args:
+            endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000).
+            access_key: MinIO access key ID for authentication.
+            secret_key: MinIO secret access key for authentication.
+            region: MinIO region name (e.g., us-east-1).
+            use_ssl: Whether to use SSL/TLS for connections.
+            max_retry_attempts: Maximum number of retry attempts (e.g., 3).
+            retry_mode: Retry policy to apply (standard, legacy, adaptive).
+            connect_timeout: Connection timeout in seconds.
+            read_timeout: Read timeout in seconds.
+            logger: Logger used for observability.
+        """
+        self.endpoint_url = endpoint_url
+        self.access_key = access_key
+        self.secret_key = secret_key
+        self.region = region
+        self.use_ssl = use_ssl
+        self.max_retry_attempts = max_retry_attempts
+        self.retry_mode = retry_mode
+        self.connect_timeout = connect_timeout
+        self.read_timeout = read_timeout
+        self.logger = logger
+
+        boto_config = Config(
+            region_name=region,
+            retries={
+                'max_attempts': max_retry_attempts,
+                'mode': retry_mode,
+            },
+            connect_timeout=connect_timeout,
+            read_timeout=read_timeout,
+        )
+
+        self.minio_client = boto3.client(
+            's3',
+            endpoint_url=endpoint_url,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            config=boto_config,
+            use_ssl=use_ssl,
+        )
+
+        self.logger.info(f'MinIO client initialized successfully: {endpoint_url}')
+
+    def fetch_file(self, bucket_name: str, file_name: str) -> BytesIO:
+        """
+        Fetch an object from MinIO and return its contents as `BytesIO`.
+
+        Args:
+            bucket_name: MinIO bucket where the object resides.
+            file_name: Object key to download inside the bucket.
+
+        Returns:
+            BytesIO: File-like stream containing the downloaded bytes.
+
+        Raises:
+            OSError: If the download fails (network, permissions, missing key, etc.).
+        """
+        response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
+
+        with response['Body'] as body:
+            file_content = body.read()
+
+        file_size = len(file_content)
+
+        self.logger.info(
+            f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)'
+        )
+
+        return BytesIO(file_content)
+
+    def delete_file(self, bucket_name: str, file_name: str) -> None:
+        """
+        Remove an object from MinIO storage.
+
+        Args:
+            bucket_name: Bucket that contains the object.
+            file_name: Object key to delete.
+        """
+        self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
+        self.logger.info(f'File deleted successfully: {bucket_name}/{file_name}')

model_manager/utils/repository/training_repository.py

@@ -66,19 +66,14 @@ class TrainingRepository:
             ValueError: If transformed data is empty
             Exception: If data loading, preprocessing, or training fails
         """
-        # Load data from BytesIO
         data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
-
-        # Initialize and fit data preprocessor
-        process_data = self.init_data_preprocessor(params)
+        process_data = self._init_data_preprocessor(params)
         process_data.fit(data)
         data_view = process_data.transform(data)
 
-        # Validate transformed data
         if len(data_view) <= 0:
             raise ValueError('Data view is empty after transformation')
 
-        # Split data into train/test sets
         x_train, x_test, y_train, y_test = split_train_test(
             data_view[params.variable_columns],
             data_view[params.target_variable],
@@ -87,18 +82,19 @@ class TrainingRepository:
             random_state=42,
         )
 
-        # Prepare training data
         data_train = pd.concat([x_train, y_train], axis=1)
-        scaler_dict = self.init_scaler_dict(process_data, params)
+        scaler_dict = self._init_scaler_dict(process_data, params)
 
-        # Create and train linear regression model
         regr = LinearRegressionModel(
             target_variable=params.target_variable,
             variable_columns=params.variable_columns,
         )
-        regr.fit(data_train)
 
-        # Return training result
+        regr.fit(data_train)
+        self.logger.info(
+            f'Model trained successfully - experiment run id: {params.experiment_run_id}'
+        )
+
         return TrainModelResult(
             params=params,
             process_data=process_data,
@@ -110,7 +106,85 @@ class TrainingRepository:
             scaler_dict=scaler_dict,
         )
 
-    def init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
+    def after_train_calculation(
+        self, params: TrainModelParams, tmr: TrainModelResult
+    ) -> TrainModelResult:
+        """
+        Perform post-training calculations: predictions, denormalization, and metrics.
+
+        This method completes the training pipeline by:
+        1. Making predictions on test set
+        2. Denormalizing all data (if scaler was used)
+        3. Reordering data by index
+        4. Calculating evaluation metrics (MSE, MAE, R²)
+
+        Args:
+            params: Training parameters used during model training
+            tmr: Result object from training
+
+        Returns:
+            TrainModelResult: Updated result with predictions, denormalized data,
+                            and metrics (mse_val, mae_val, r2_val)
+        """
+        y_pred_array = tmr.regr.predict(tmr.x_test)
+
+        if params.use_scaler:
+            scaler = tmr.process_data.get_scaler()
+
+            # If using custom scaler with denormalize_* helpers
+            if hasattr(scaler, 'denormalize_single_input'):
+                for col in params.variable_columns:
+                    tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
+                    tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
+
+                tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
+                tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
+                y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
+            else:
+                # Fallback for sklearn StandardScaler: only inverse-transform features
+                feature_cols = getattr(
+                    tmr.process_data, 'feature_names_order', params.variable_columns
+                )
+                # Ensure columns are in the same order used during fit
+                x_train_features = tmr.x_train[feature_cols]
+                x_test_features = tmr.x_test[feature_cols]
+
+                tmr.x_train[feature_cols] = scaler.inverse_transform(x_train_features)
+                tmr.x_test[feature_cols] = scaler.inverse_transform(x_test_features)
+                # Target was not scaled with StandardScaler in preprocessing; leave y as-is
+
+        tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
+        tmr.y_pred.name = f'{params.target_variable}_pred'
+
+        tmr.x_train = tmr.x_train.sort_index()
+        tmr.x_test = tmr.x_test.sort_index()
+        tmr.y_train = tmr.y_train.sort_index()
+        tmr.y_test = tmr.y_test.sort_index()
+        tmr.y_pred = tmr.y_pred.sort_index()
+
+        assert tmr.y_pred is not None, 'y_pred should be set at this point'
+
+        tmr.mse_val = round(
+            mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
+            2,
+        )
+
+        tmr.mae_val = round(
+            mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
+            2,
+        )
+
+        tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
+
+        # Extract model equation
+        tmr.equation = self._extract_model_equation(tmr.regr, params)
+
+        self.logger.info(
+            f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
+        )
+        return tmr
+
+    def _init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
         """
         Initialize dictionary containing scaling parameters for features and target.
 
@@ -152,69 +226,7 @@ class TrainingRepository:
 
         return scaler_dict
 
-    def after_train_calculation(
-        self, params: TrainModelParams, tmr: TrainModelResult
-    ) -> TrainModelResult:
-        """
-        Perform post-training calculations: predictions, denormalization, and metrics.
-
-        This method completes the training pipeline by:
-        1. Making predictions on test set
-        2. Denormalizing all data (if scaler was used)
-        3. Reordering data by index
-        4. Calculating evaluation metrics (MSE, MAE, R²)
-
-        Args:
-            params: Training parameters used during model training
-            tmr: Result object from training
-
-        Returns:
-            TrainModelResult: Updated result with predictions, denormalized data,
-                            and metrics (mse_val, mae_val, r2_val)
-        """
-        # Make predictions on test set
-        y_pred_array = tmr.regr.predict(tmr.x_test)
-
-        # Denormalize data if scaler was used
-        if params.use_scaler:
-            scaler = tmr.process_data.get_scaler()
-
-            # Denormalize features
-            for col in params.variable_columns:
-                tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
-                tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
-
-            # Denormalize target variable
-            tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
-            tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
-            y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
-
-        # Add index to predictions
-        tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
-        tmr.y_pred.name = f'{params.target_variable}_pred'
-
-        # Reorder all data by index
-        tmr.x_train = tmr.x_train.sort_index()
-        tmr.x_test = tmr.x_test.sort_index()
-        tmr.y_train = tmr.y_train.sort_index()
-        tmr.y_test = tmr.y_test.sort_index()
-        tmr.y_pred = tmr.y_pred.sort_index()
-
-        # Calculate evaluation metrics
-        assert tmr.y_pred is not None, 'y_pred should be set at this point'
-        tmr.mse_val = round(
-            mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
-            2,
-        )
-        tmr.mae_val = round(
-            mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
-            2,
-        )
-        tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
-
-        return tmr
-
-    def init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
+    def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
         """
         Initialize DataPreprocessor with training parameters.
 
@@ -238,5 +250,55 @@ class TrainingRepository:
             upp_lim=params.upp_lim,
             window=params.window,
             scaler_name='Standard Scaler' if params.use_scaler else 'None',
+            scaler_params={} if params.use_scaler else None,
             ar_var=params.target_variable if params.include_ar else None,
         )
+
+    def _extract_model_equation(
+        self, regr: LinearRegressionModel, params: TrainModelParams
+    ) -> dict:
+        """
+        Extract the linear regression equation coefficients and create equation metadata.
+
+        This method extracts the coefficients and intercept from the trained model
+        and creates a structured dictionary containing the equation information
+        for serialization as JSON artifact.
+
+        Args:
+            regr: Trained LinearRegressionModel object
+            params: Training parameters containing variable information
+
+        Returns:
+            dict: Equation metadata containing:
+                - target_variable: Name of the target variable
+                - coefficients: Dictionary mapping variable names to coefficients
+                - intercept: Model intercept value
+                - equation_string: Human-readable equation string
+                - latex_equation: LaTeX formatted equation
+        """
+        coefficients = regr.regr.coef_
+        intercept = regr.regr.intercept_
+
+        # Create coefficients dictionary
+        coefficients_dict = {}
+        for i, var in enumerate(params.variable_columns):
+            coefficients_dict[var] = float(coefficients[i])
+
+        # Create equation string
+        equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
+        equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(
+            equation_parts
+        )
+
+        # Create LaTeX equation
+        latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()]
+        latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts)
+
+        return {
+            'target_variable': params.target_variable,
+            'coefficients': coefficients_dict,
+            'intercept': float(intercept),
+            'equation_string': equation_string,
+            'latex_equation': latex_equation,
+            'model_type': 'Linear Regression',
+        }

model_manager/worker/worker.py

@@ -33,7 +33,6 @@ with workflow.unsafe.imports_passed_through():
 
     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 model_manager.activities.activities import Activities
@@ -43,6 +42,7 @@ with workflow.unsafe.imports_passed_through():
         build_mongodb_config,
         build_postgres_config,
     )
+    from model_manager.utils.logger_helper import get_logger
     from model_manager.workflows.train_model import TrainModel
 
 POD_ID = os.getenv('POD_ID')
@@ -73,10 +73,7 @@ async def main():
 
     metadata = {
         'pod_id': POD_ID,
-        'model_name': '-',
-        'model_id': '-',
-        'workflow_name': '-',
-        'schedule_name': '-',
+        'workflow_name': 'train_model',
     }
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
@@ -128,18 +125,10 @@ async def main():
             task_queue='train_model-queue',
             workflows=[TrainModel],
             activities=[
-                # Training & Validation
+                activities.update_experiment_run,
                 activities.validate_train_params,
                 activities.train_model,
-                # MLFlow
-                activities.save_model,
-                # MinIO
-                activities.fetch_file_from_minio,
-                activities.delete_file_from_minio,
-                # Filesystem
-                activities.cleanup_run_directory,
-                # Database
-                activities.update_experiment_run,
+                activities.cleanup_resources,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,
@@ -163,10 +152,8 @@ async def main():
     except BaseException as e:  # noqa: BLE001
         logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
     finally:
-        if notification_handler:
-            notification_handler.shutdown()
-        if activities:
-            await activities.shutdown()
+        notification_handler.shutdown()
+        await activities.shutdown()
         # Exit with a non-zero status code to indicate failure to Kubernetes
         metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
         sys.exit(1)

model_manager/workflows/train_model.py

@@ -20,18 +20,15 @@ with workflow.unsafe.imports_passed_through():
 
     from model_manager.activities.activities import Activities
     from model_manager.activities.experiment_tracking import UpdateType
+    from model_manager.utils.exceptions import ModelTrainingError
     from model_manager.utils.models.experiment_status import ExperimentStatus
     from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
 
     # Activity Timeouts (in seconds) - Configurable via environment variables
     # Defaults are designed to handle large files (up to 200MB)
     TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
-    TIMEOUT_DOWNLOAD_FILE = int(os.getenv('TIMEOUT_DOWNLOAD_FILE', '600'))
-    TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '1800'))
-    TIMEOUT_SAVE_MODEL = int(os.getenv('TIMEOUT_SAVE_MODEL', '300'))
-    TIMEOUT_CLEANUP_DIRECTORY = int(os.getenv('TIMEOUT_CLEANUP_DIRECTORY', '60'))
-    TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '60'))
+    TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '2700'))
+    TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '120'))
     TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
 
     # Retry Policies - Granular strategies for different operation types
@@ -48,14 +45,6 @@ with workflow.unsafe.imports_passed_through():
         maximum_attempts=1,
     )
 
-    # Moderate retry with backoff for MLFlow operations
-    mlflow_retry_policy = RetryPolicy(
-        initial_interval=timedelta(seconds=5),
-        maximum_interval=timedelta(seconds=30),
-        backoff_coefficient=2.0,
-        maximum_attempts=3,
-    )
-
     # Database retry with exponential backoff
     database_retry_policy = RetryPolicy(
         initial_interval=timedelta(seconds=2),
@@ -64,13 +53,7 @@ with workflow.unsafe.imports_passed_through():
         maximum_attempts=5,
     )
 
-    # Filesystem retry for cleanup operations
-    filesystem_retry_policy = RetryPolicy(
-        initial_interval=timedelta(seconds=2),
-        maximum_interval=timedelta(seconds=10),
-        backoff_coefficient=1.5,
-        maximum_attempts=3,
-    )
+POD_ID = os.getenv('POD_ID')
 
 
 @workflow.defn(name='train_model')
@@ -113,44 +96,31 @@ class TrainModel:
         Raises:
             ValueError: If experiment_run_id is missing or invalid
         """
-        # CRITICAL: Validate experiment_run_id first
-        # Without it, we cannot update database status, so fail immediately
         experiment_run_id = self._validate_experiment_run_id(input_data)
 
         metadata = {
             'metadata': {
+                'pod_id': POD_ID,
                 'experiment_run_id': experiment_run_id,
                 'workflow_name': 'train_model',
             }
         }
 
-        # Step 1: Validate and convert training parameters
         train_params = await self._validate_training_parameters(
             input_data, experiment_run_id, metadata
         )
 
-        # Step 2 & 3: Download from MinIO and Train model
-        # The _download_and_train_model method handles both steps:
-        # - Downloads file from MinIO (returns BytesIO)
-        # - Trains model with the downloaded file
-        # Any error (download OR training) = TRAINING_ERROR
-        train_result = await self._download_and_train_model(
+        train_result = await self._train_model(
             train_params=train_params,
             experiment_run_id=experiment_run_id,
             metadata=metadata,
         )
 
-        # Step 4: Save model to MLFlow
-        saved_result = await self._save_model_to_mlflow(
-            train_result=train_result,
-            experiment_run_id=experiment_run_id,
-            metadata=metadata,
-        )
-
-        # Step 5: Cleanup resources and delete file from MinIO
         await self._cleanup_resources(
-            saved_result=saved_result,
             experiment_run_id=experiment_run_id,
+            run_dir=(train_result.get('run_dir') or ''),
+            bucket_name=train_params.bucket_name,
+            file_name=train_params.file_name,
             metadata=metadata,
         )
 
@@ -174,16 +144,12 @@ class TrainModel:
         experiment_run_id = input_data.get('experiment_run_id')
 
         if experiment_run_id is None:
-            error_msg = 'experiment_run_id is required but was not provided'
-            workflow.logger.error(error_msg)
-            raise ValueError(error_msg)
+            raise ValueError('experiment_run_id is required but was not provided')
 
         if not isinstance(experiment_run_id, int):
-            error_msg = (
+            raise ValueError(
                 f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}'
             )
-            workflow.logger.error(error_msg)
-            raise ValueError(error_msg)
 
         return experiment_run_id
 
@@ -211,21 +177,17 @@ class TrainModel:
         Raises:
             Exception: If validation fails (after updating DB status)
         """
-        validation_input = {
-            **metadata,
-            **input_data,  # All training params at same level
-        }
-
         try:
-            # Execute validation activity
             train_params = await workflow.execute_activity_method(
                 Activities.validate_train_params,
-                validation_input,
-                retry_policy=no_retry_policy,  # Validation errors are permanent
+                {
+                    **metadata,
+                    **input_data,
+                },
+                retry_policy=no_retry_policy,
                 start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
             )
 
-            # Validation succeeded: Update status
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
@@ -234,41 +196,23 @@ class TrainModel:
             )
 
             return train_params
-
         except Exception as e:
-            # Validation failed: Update status with error
-            error_message = str(e)
-            error_type = type(e).__name__
-
-            # Log with rich context for debugging
-            workflow.logger.error(
-                f'[VALIDATION_ERROR] Experiment {experiment_run_id} validation failed',
-                extra={
-                    'step': 'validate_training_parameters',
-                    'experiment_run_id': experiment_run_id,
-                    'error_type': error_type,
-                    'error_message': error_message,
-                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
-                },
-            )
-
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
                 update_type=UpdateType.STATUS_WITH_ERROR,
                 status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
-                error_message=error_message,
+                error_message=self._extract_error_message(e),
             )
 
-            # Re-raise exception to stop workflow
             raise
 
-    async def _download_and_train_model(
+    async def _train_model(
         self,
         train_params: TrainModelParams,
         experiment_run_id: int,
         metadata: dict[str, Any],
-    ) -> TrainModelResult:
+    ) -> dict[str, str | None]:
         """
         Download file from MinIO and train model.
 
@@ -287,167 +231,53 @@ class TrainModel:
         Raises:
             Exception: If download or training fails (after updating DB status)
         """
-        uploaded_file = None
-
         try:
-            # Step 1: Download file from MinIO
-            download_input = {
-                **metadata,
-                'bucket_name': train_params.bucket_name,
-                'file_name': train_params.file_name,
-            }
-
-            uploaded_file = await workflow.execute_activity_method(
-                Activities.fetch_file_from_minio,
-                download_input,
-                retry_policy=network_retry_policy,  # Fast retry for network issues
-                start_to_close_timeout=timedelta(seconds=TIMEOUT_DOWNLOAD_FILE),
-            )
-
-            # Step 2: Train model with downloaded file
-            train_input = {
-                **metadata,
-                'uploaded_file': uploaded_file,
-                'train_params': train_params,
-            }
-
             train_result = await workflow.execute_activity_method(
                 Activities.train_model,
-                train_input,
-                retry_policy=no_retry_policy,  # Training errors are permanent (bad data)
+                {
+                    **metadata,
+                    'train_params': train_params,
+                },
+                retry_policy=no_retry_policy,
                 start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
             )
 
-            # Training succeeded: Update status
-            await self._update_experiment_run(
-                metadata=metadata,
-                experiment_run_id=experiment_run_id,
-                update_type=UpdateType.STATUS,
-                status=ExperimentStatus.TRAINING_SUCCESS,
-            )
-
-            return train_result
-
-        except Exception as e:
-            # Download or training failed: Update status with error
-            error_message = str(e)
-            error_type = type(e).__name__
-
-            # Log with rich context for debugging
-            workflow.logger.error(
-                f'[TRAINING_ERROR] Experiment {experiment_run_id} training failed',
-                extra={
-                    'step': 'download_and_train_model',
-                    'experiment_run_id': experiment_run_id,
-                    'experiment_name': train_params.experiment_name,
-                    'bucket_name': train_params.bucket_name,
-                    'file_name': train_params.file_name,
-                    'error_type': error_type,
-                    'error_message': error_message,
-                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
-                },
-            )
-
-            await self._update_experiment_run(
-                metadata=metadata,
-                experiment_run_id=experiment_run_id,
-                update_type=UpdateType.STATUS_WITH_ERROR,
-                status=ExperimentStatus.TRAINING_ERROR,
-                error_message=error_message,
-            )
-
-            # Re-raise exception to stop workflow
-            raise
-
-        finally:
-            # Ensure BytesIO is closed
-            if uploaded_file and hasattr(uploaded_file, 'close'):
-                uploaded_file.close()
-
-    async def _save_model_to_mlflow(
-        self,
-        train_result: TrainModelResult,
-        experiment_run_id: int,
-        metadata: dict[str, Any],
-    ) -> TrainModelResult:
-        """
-        Save trained model to MLFlow.
-
-        This method calls the save_model activity to save the trained model and
-        its artifacts to MLFlow. On success, updates DB status to MLFLOW_SENT
-        with MODEL_SAVED type and run_name. On error, updates DB status to
-        MLFLOW_SEND_ERROR.
-
-        Args:
-            train_result: TrainModelResult from training step
-            experiment_run_id: Validated experiment run ID
-            metadata: Workflow execution metadata
-
-        Returns:
-            TrainModelResult: Updated training result with MLFlow run name
-
-        Raises:
-            Exception: If model saving fails (after updating DB status)
-        """
-        try:
-            # Save model to MLFlow
-            save_input = {
-                **metadata,
-                'train_result': train_result,
-            }
-
-            saved_result = await workflow.execute_activity_method(
-                Activities.save_model,
-                save_input,
-                retry_policy=mlflow_retry_policy,  # Retry MLFlow with backoff
-                start_to_close_timeout=timedelta(seconds=TIMEOUT_SAVE_MODEL),
-            )
-
-            # Model saved successfully: Update status to MLFLOW_SENT with run_name
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
                 update_type=UpdateType.MODEL_SAVED,
                 status=ExperimentStatus.MLFLOW_SENT,
-                run_name=saved_result.run_name,
+                run_name=train_result.get('run_name'),
             )
 
-            return saved_result
-
+            return train_result
         except Exception as e:
-            # Model saving failed: Update status with error
-            error_message = str(e)
-            error_type = type(e).__name__
+            # Mapear flags -> status
+            # False/False: erro no treino
+            # True/False: erro ao salvar (MLflow)
+            # False/True: estado inconsistente, tratar como erro de treino
+            # True/True: não deveria cair aqui; tratar como erro genérico de treino
+            status = ExperimentStatus.TRAINING_ERROR
 
-            # Log with rich context for debugging
-            workflow.logger.error(
-                f'[MLFLOW_ERROR] Experiment {experiment_run_id} model save failed',
-                extra={
-                    'step': 'save_model_to_mlflow',
-                    'experiment_run_id': experiment_run_id,
-                    'experiment_name': train_result.params.experiment_name,
-                    'run_dir': train_result.run_dir if hasattr(train_result, 'run_dir') else None,
-                    'error_type': error_type,
-                    'error_message': error_message,
-                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
-                },
-            )
+            if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
+                status = ExperimentStatus.MLFLOW_SEND_ERROR
 
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
                 update_type=UpdateType.STATUS_WITH_ERROR,
-                status=ExperimentStatus.MLFLOW_SEND_ERROR,
-                error_message=error_message,
+                status=status,
+                error_message=self._extract_error_message(e),
             )
 
-            # Re-raise exception to stop workflow
             raise
 
     async def _cleanup_resources(
         self,
-        saved_result: TrainModelResult,
         experiment_run_id: int,
+        run_dir: str,
+        bucket_name: str,
+        file_name: str,
         metadata: dict[str, Any],
     ) -> None:
         """
@@ -466,75 +296,33 @@ class TrainModel:
             Exception: If cleanup fails (after updating DB status)
         """
         try:
-            # Step 1: Remove temporary run directory via activity (deterministic)
-            if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
-                cleanup_input = {
-                    **metadata,
-                    'run_dir': saved_result.run_dir,
-                }
-
-                await workflow.execute_activity_method(
-                    Activities.cleanup_run_directory,
-                    cleanup_input,
-                    retry_policy=filesystem_retry_policy,  # Retry filesystem operations
-                    start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_DIRECTORY),
-                )
-
-                workflow.logger.info(
-                    f'Run directory cleanup completed for experiment {experiment_run_id}'
-                )
-
-            # Step 2: Delete file from MinIO
-            delete_input = {
-                **metadata,
-                'bucket_name': saved_result.params.bucket_name,
-                'file_name': saved_result.params.file_name,
-            }
-
             await workflow.execute_activity_method(
-                Activities.delete_file_from_minio,
-                delete_input,
-                retry_policy=network_retry_policy,  # Fast retry for network issues
+                Activities.cleanup_resources,
+                {
+                    **metadata,
+                    'run_dir': run_dir,
+                    'bucket_name': bucket_name,
+                    'file_name': file_name,
+                },
+                retry_policy=network_retry_policy,
                 start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
             )
 
-            # Cleanup succeeded: Update status to FILE_DELETED
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
                 update_type=UpdateType.STATUS,
                 status=ExperimentStatus.FILE_DELETED,
             )
-
         except Exception as e:
-            # Cleanup failed: Update status with error
-            error_message = str(e)
-            error_type = type(e).__name__
-
-            # Log with rich context for debugging
-            workflow.logger.error(
-                f'[CLEANUP_ERROR] Experiment {experiment_run_id} cleanup failed',
-                extra={
-                    'step': 'cleanup_resources',
-                    'experiment_run_id': experiment_run_id,
-                    'bucket_name': saved_result.params.bucket_name,
-                    'file_name': saved_result.params.file_name,
-                    'run_dir': saved_result.run_dir if hasattr(saved_result, 'run_dir') else None,
-                    'error_type': error_type,
-                    'error_message': error_message,
-                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
-                },
-            )
-
             await self._update_experiment_run(
                 metadata=metadata,
                 experiment_run_id=experiment_run_id,
                 update_type=UpdateType.STATUS_WITH_ERROR,
                 status=ExperimentStatus.FILE_DELETE_ERROR,
-                error_message=error_message,
+                error_message=self._extract_error_message(e),
             )
 
-            # Re-raise exception to stop workflow
             raise
 
     async def _update_experiment_run(
@@ -567,17 +355,34 @@ class TrainModel:
             'status': status,
         }
 
-        # Add error_message only if provided
         if error_message is not None:
             update_input['error_message'] = error_message
 
-        # Add run_name only if provided
         if run_name is not None:
             update_input['run_name'] = run_name
 
         await workflow.execute_activity_method(
             Activities.update_experiment_run,
             update_input,
-            retry_policy=database_retry_policy,  # Retry DB with exponential backoff
+            retry_policy=database_retry_policy,
             start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
         )
+
+    def _extract_error_message(self, exc: Exception) -> str:
+        message_parts: list[str] = []
+        seen: set[int] = set()
+        current: Exception | None = exc
+
+        while current and id(current) not in seen:
+            seen.add(id(current))
+            text = str(current).strip()
+
+            if text and text not in message_parts:
+                message_parts.append(text)
+
+            current = getattr(current, 'cause', None)
+
+        if not message_parts:
+            return repr(exc)
+
+        return ' | '.join(message_parts)

run_coverage.sh

@@ -1,11 +0,0 @@
-#!/bin/bash
-
-# Exit on any error
-set -e
-
-echo "Activating virtual environment..."
-source ./venv/bin/activate
-
-pytest --cov=model_manager --cov-report=html
-
-xdg-open htmlcov/index.html

scripts/run_training_test.py

@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""Utility script to trigger the training workflow end-to-end for testing.
+
+Steps performed:
+1. Upload the CSV test dataset to MinIO using the configured `mc` alias.
+2. Insert a new experiment_run record in Postgres and capture the generated ID.
+3. Trigger the Temporal `train_model` workflow with the correct payload.
+
+Prerequisites:
+- `mc` CLI configured with alias defined in MINIO_ALIAS.
+- PostgreSQL accessible with credentials in environment variables or defaults.
+- Temporal server reachable without TLS on TEMPORAL_HOST / TEMPORAL_NAMESPACE.
+- Python dependencies installed (see requirements.txt / requirements-dev.txt).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import subprocess
+import sys
+import uuid
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import psycopg2
+from psycopg2.extras import Json
+from temporalio import client
+
+DOCS_PATH = Path('docs/test-model-data.csv')
+MINIO_ALIAS = os.getenv('MINIO_ALIAS', 'suse')
+MINIO_BUCKET = os.getenv('MINIO_BUCKET', 'model-training')
+
+POSTGRES_CONFIG = {
+    'host': os.getenv('POSTGRES_HOST', 'localhost'),
+    'port': os.getenv('POSTGRES_PORT', '55432'),
+    'user': os.getenv('POSTGRES_USER', 'postgres'),
+    'password': os.getenv(
+        'POSTGRES_PASSWORD',
+        'nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3',
+    ),
+    'dbname': os.getenv('POSTGRES_DBNAME', 'sientia-core-mlops-bff'),
+}
+
+TEMPORAL_HOST = os.getenv('TEMPORAL_HOST', 'localhost:37463')
+TEMPORAL_NAMESPACE = os.getenv('TEMPORAL_NAMESPACE', 'model-manager')
+TEMPORAL_TASK_QUEUE = os.getenv('TEMPORAL_TASK_QUEUE', 'train_model-queue')
+TEMPORAL_WORKFLOW = os.getenv('TEMPORAL_WORKFLOW', 'train_model')
+
+BASE_REQUEST_DATA = {
+    'experimentName': 'model-manager-test-01',
+    'username': 'bruno.domingues@aignosi.com.br',
+    'modelType': 'Linear Regression',
+    'targetVariable': '03CV020/CORRENTE_N_M1_PV(Value)',
+    'variableColumns': ['303-WIT-200(Value)'],
+    'lagTrain': 0,
+    'lagVal': 0,
+    'remStaticWin': False,
+    'lowLim': {},
+    'uppLim': {},
+    'window': 0,
+    'useScaler': False,
+    'includeAr': False,
+    'trainSize': 80,
+    'shuffle': True,
+    'lineSeparator': ',',
+    'decimalSeparator': '.',
+    'removedIntervals': [],
+}
+
+
+def _ensure_source_file(path: Path) -> None:
+    if not path.exists():
+        raise FileNotFoundError(f'Test dataset not found at {path.resolve()}')
+
+
+def upload_to_minio(source_path: Path) -> str:
+    """Upload the CSV to MinIO using the mc CLI and return the object name."""
+    _ensure_source_file(source_path)
+    timestamp = datetime.utcnow().strftime('%Y%m%d-%H%M%S')
+    object_name = f'test-model-data-{timestamp}.csv'
+    target_uri = f'{MINIO_ALIAS}/{MINIO_BUCKET}/{object_name}'
+
+    subprocess.run(  # noqa: S603
+        ['mc', 'cp', str(source_path), target_uri],  # noqa: S607
+        check=True,
+    )
+    return object_name
+
+
+def insert_experiment_run(file_name: str, request_data: dict) -> int:
+    """Insert experiment_run record and return the generated ID."""
+    now = datetime.utcnow()
+    payload = {
+        **request_data,
+        'fileName': file_name,
+        'bucketName': MINIO_BUCKET,
+    }
+
+    insert_sql = """
+        INSERT INTO experiment_run (
+            experiment_name,
+            username,
+            status,
+            created_at,
+            updated_at,
+            bucket_name,
+            file_name,
+            request_data
+        )
+        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+        RETURNING id;
+    """
+
+    with psycopg2.connect(**POSTGRES_CONFIG) as conn:
+        with conn.cursor() as cur:
+            cur.execute(
+                insert_sql,
+                (
+                    request_data['experimentName'],
+                    request_data['username'],
+                    'MAGE_REQUEST_SENT',
+                    now,
+                    now,
+                    MINIO_BUCKET,
+                    file_name,
+                    Json(payload),
+                ),
+            )
+            experiment_run_id = cur.fetchone()[0]
+
+    return experiment_run_id
+
+
+def build_workflow_payload(
+    experiment_run_id: int,
+    file_name: str,
+    request_data: dict,
+) -> dict:
+    """Convert camelCase request data to snake_case and enrich with runtime values."""
+    return {
+        'experiment_run_id': experiment_run_id,
+        'experiment_name': request_data['experimentName'],
+        'username': request_data['username'],
+        'model_type': request_data['modelType'],
+        'target_variable': request_data['targetVariable'],
+        'variable_columns': request_data['variableColumns'],
+        'lag_train': request_data['lagTrain'],
+        'lag_val': request_data['lagVal'],
+        'rem_static_win': request_data['remStaticWin'],
+        'low_lim': request_data['lowLim'],
+        'upp_lim': request_data['uppLim'],
+        'window': request_data['window'],
+        'use_scaler': request_data['useScaler'],
+        'include_ar': request_data['includeAr'],
+        'train_size': request_data['trainSize'],
+        'shuffle': request_data['shuffle'],
+        'bucket_name': MINIO_BUCKET,
+        'file_name': file_name,
+        'line_separator': request_data['lineSeparator'],
+        'decimal_separator': request_data['decimalSeparator'],
+        'removed_intervals': request_data['removedIntervals'],
+    }
+
+
+async def trigger_temporal_workflow(workflow_input: dict) -> str:
+    """Connect to Temporal and trigger the training workflow."""
+    temporal_client = await client.Client.connect(
+        target_host=TEMPORAL_HOST,
+        namespace=TEMPORAL_NAMESPACE,
+    )
+
+    workflow_id = f'train-model-test-{uuid.uuid4()}'
+    await temporal_client.execute_workflow(
+        TEMPORAL_WORKFLOW,
+        workflow_input,
+        id=workflow_id,
+        task_queue=TEMPORAL_TASK_QUEUE,
+        execution_timeout=timedelta(minutes=5),
+        run_timeout=timedelta(minutes=5),
+        task_timeout=timedelta(minutes=5),
+    )
+    return workflow_id
+
+
+def main() -> None:
+    try:
+        uploaded_file_name = upload_to_minio(DOCS_PATH)
+    except subprocess.CalledProcessError as exc:
+        print(f'Failed to upload file to MinIO: {exc}', file=sys.stderr)
+        sys.exit(1)
+    except FileNotFoundError as exc:
+        print(str(exc), file=sys.stderr)
+        sys.exit(1)
+
+    experiment_request = BASE_REQUEST_DATA.copy()
+
+    try:
+        experiment_run_id = insert_experiment_run(uploaded_file_name, experiment_request)
+    except psycopg2.Error as exc:
+        print(f'Database error while inserting experiment_run: {exc}', file=sys.stderr)
+        sys.exit(1)
+
+    workflow_payload = build_workflow_payload(
+        experiment_run_id=experiment_run_id,
+        file_name=uploaded_file_name,
+        request_data=experiment_request,
+    )
+
+    try:
+        workflow_id = asyncio.run(trigger_temporal_workflow(workflow_payload))
+    except Exception as exc:  # noqa: BLE001
+        print(f'Failed to start Temporal workflow: {exc}', file=sys.stderr)
+        sys.exit(1)
+
+    print(
+        json.dumps(
+            {
+                'experiment_run_id': experiment_run_id,
+                's3_object_name': uploaded_file_name,
+                'workflow_id': workflow_id,
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == '__main__':
+    main()

setup_port_forwards.sh

@@ -1,111 +0,0 @@
-#!/bin/bash
-
-# Exit on any error
-set -e
-
-# Color codes for output
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-NC='\033[0m' # No Color
-
-echo -e "${GREEN}=== Port Forward Setup Script ===${NC}\n"
-
-# Define port forwards: LOCAL_PORT:NAMESPACE:SERVICE:REMOTE_PORT:DESCRIPTION
-PORT_FORWARDS=(
-    "5432:paradedb:paradedb-rw:5432:PostgreSQL"
-    "45249:sientia-tracker:sientia-tracker-mlflow-tracking:80:MLflow"
-    "37463:temporal:temporal-frontend:7233:Temporal"
-    "8080:temporal:temporal-web:8080:Temporal UI"
-    "42297:mongodb:my-release-mongodb:27017:MongoDB"
-    "36577:minio:minio:9000:MinIO"
-)
-
-# Step 1: Kill existing port-forward jobs for these services
-echo -e "${YELLOW}Step 1: Checking for existing port-forward jobs...${NC}"
-
-for pf in "${PORT_FORWARDS[@]}"; do
-    IFS=':' read -r local_port namespace service remote_port description <<< "$pf"
-    
-    # Check if there's a job with this service name
-    existing_jobs=$(jobs -l | grep "kubectl.*port-forward.*svc/$service" || true)
-    
-    if [ -n "$existing_jobs" ]; then
-        echo -e "${YELLOW}  Found existing port-forward for $description ($service)${NC}"
-        # Extract PIDs and kill them
-        pids=$(echo "$existing_jobs" | awk '{print $2}')
-        for pid in $pids; do
-            echo -e "${YELLOW}  Killing job with PID $pid${NC}"
-            kill "$pid" 2>/dev/null || true
-        done
-    fi
-done
-
-# Wait a moment for ports to be released
-sleep 1
-
-echo -e "${GREEN}  Cleanup complete${NC}\n"
-
-# Step 2: Check if any of the ports are already in use
-echo -e "${YELLOW}Step 2: Checking if ports are available...${NC}"
-
-ports_in_use=()
-
-for pf in "${PORT_FORWARDS[@]}"; do
-    IFS=':' read -r local_port namespace service remote_port description <<< "$pf"
-    
-    # Check if port is in use using lsof or netstat
-    if command -v lsof &> /dev/null; then
-        if lsof -Pi :$local_port -sTCP:LISTEN -t >/dev/null 2>&1; then
-            ports_in_use+=("$local_port:$description")
-        fi
-    elif command -v netstat &> /dev/null; then
-        if netstat -tuln | grep -q ":$local_port "; then
-            ports_in_use+=("$local_port:$description")
-        fi
-    elif command -v ss &> /dev/null; then
-        if ss -tuln | grep -q ":$local_port "; then
-            ports_in_use+=("$local_port:$description")
-        fi
-    fi
-done
-
-# If any ports are in use, report and exit
-if [ ${#ports_in_use[@]} -gt 0 ]; then
-    echo -e "${RED}ERROR: The following ports are already in use:${NC}"
-    for port_info in "${ports_in_use[@]}"; do
-        IFS=':' read -r port desc <<< "$port_info"
-        echo -e "${RED}  - Port $port (for $desc)${NC}"
-    done
-    echo -e "\n${RED}Please free these ports before running this script.${NC}"
-    exit 1
-fi
-
-echo -e "${GREEN}  All ports are available${NC}\n"
-
-# Step 3: Create all port forwards
-echo -e "${YELLOW}Step 3: Creating port forwards...${NC}"
-
-for pf in "${PORT_FORWARDS[@]}"; do
-    IFS=':' read -r local_port namespace service remote_port description <<< "$pf"
-    
-    echo -e "${GREEN}  Starting port-forward: $description${NC}"
-    echo -e "    Local port: $local_port -> $namespace/$service:$remote_port"
-    
-    kubectl -n "$namespace" port-forward "svc/$service" "$local_port:$remote_port" &
-    
-    # Give it a moment to start
-    sleep 0.5
-done
-
-echo -e "\n${GREEN}=== All port forwards created successfully ===${NC}"
-echo -e "\n${YELLOW}Active port forwards:${NC}"
-for pf in "${PORT_FORWARDS[@]}"; do
-    IFS=':' read -r local_port namespace service remote_port description <<< "$pf"
-    echo -e "  - ${GREEN}localhost:$local_port${NC} -> $description ($namespace/$service)"
-done
-
-echo -e "\n${YELLOW}To stop all port forwards, run:${NC}"
-echo -e "  jobs -p | xargs kill"
-echo -e "\n${YELLOW}To view active port forwards:${NC}"
-echo -e "  jobs -l"

simulator/Dockerfile

@@ -1,30 +0,0 @@
-# syntax=docker/dockerfile:1.4
-
-FROM python:3.11-slim
-
-# Enable use of SSH agent/socket
-# This line enables SSH during build
-# (don't forget the syntax header above)
-RUN apt-get update && apt-get install -y git openssh-client && rm -rf /var/lib/apt/lists/*
-
-# Use build-time SSH mount for Git clone
-# The SSH key will NOT remain in the image
-# IMPORTANT: this block requires BuildKit
-# and the --ssh flag during docker build
-
-# SSH config to skip host key check (safe in CI/local dev)
-RUN mkdir -p /root/.ssh && echo "StrictHostKeyChecking no" > /root/.ssh/config
-
-WORKDIR /app
-
-# Clone using SSH
-ARG GIT_REPO
-ARG GIT_BRANCH=main
-
-# Mount SSH key just for this RUN
-RUN --mount=type=ssh git clone --branch ${GIT_BRANCH} ${GIT_REPO} .
-
-# Install requirements if exists
-RUN if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi
-
-CMD ["python", "server.py"]

tests/activities/test_activities.py

@@ -1,88 +1,108 @@
-from unittest.mock import ANY, MagicMock, patch
+"""Unit tests for Activities class with 100% coverage."""
 
-from pytest import mark
+import asyncio
+from unittest.mock import MagicMock, patch
 
-from model_manager.activities.activities import Activities
-from model_manager.activities.experiment_tracking import ExperimentTracking
-from model_manager.activities.mlflow import MLFlow
-from model_manager.activities.training import Training
+import pytest
 
 
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___init__(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    postgres_config = {
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_notification_handler():
+    """Create a mock notification handler."""
+    return MagicMock()
+
+
+@pytest.fixture
+def postgres_config():
+    """Create a valid PostgreSQL configuration."""
+    return {
         'host': 'localhost',
         'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
+        'user': 'testuser',
+        'password': 'testpass',
+        'dbname': 'testdb',
         'min_connections': 1,
         'max_connections': 10,
     }
 
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
 
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
+@pytest.fixture
+def mlflow_config():
+    """Create a valid MLFlow configuration."""
+    return {
+        'url': 'http://mlflow:5080',
+        'username': 'aignosi',
+        'password': 'aignosi',
+    }
+
+
+@pytest.fixture
+def minio_config():
+    """Create a valid MinIO configuration."""
+    return {
+        'endpoint_url': 'http://minio:9000',
         'access_key': 'minioadmin',
         'secret_key': 'minioadmin',
         'region': 'us-east-1',
         'use_ssl': False,
         'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
+        'retry_mode': 'standard',
+        'connect_timeout': 5,
+        'read_timeout': 5,
     }
 
-    logger = MagicMock()
-    notification_handler = MagicMock()
+
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_init_success(
+    mock_storage_repo,
+    mock_model_repo,
+    mock_training_init,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test successful initialization of Activities."""
+    from model_manager.activities.activities import Activities
 
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    assert isinstance(activities, Activities)
-    assert isinstance(activities, ExperimentTracking)
-    assert isinstance(activities, MLFlow)
-    assert isinstance(activities, Training)
+    mock_et_init.assert_called_once()
+    assert mock_et_init.call_args[1]['host'] == postgres_config['host']
+    assert mock_et_init.call_args[1]['port'] == postgres_config['port']
+    assert mock_et_init.call_args[1]['user'] == postgres_config['user']
+    assert mock_et_init.call_args[1]['password'] == postgres_config['password']
+    assert mock_et_init.call_args[1]['dbname'] == postgres_config['dbname']
+    assert mock_et_init.call_args[1]['min_connections'] == postgres_config['min_connections']
+    assert mock_et_init.call_args[1]['max_connections'] == postgres_config['max_connections']
+    assert mock_et_init.call_args[1]['logger'] is mock_logger
+    assert mock_et_init.call_args[1]['notification_handler'] is mock_notification_handler
 
-    mock_experiment_tracking_init.assert_called_once_with(
-        ANY,
-        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,
+    mock_model_repo.assert_called_once_with(
+        url=mlflow_config['url'],
+        username=mlflow_config['username'],
+        password=mlflow_config['password'],
+        logger=mock_logger,
     )
 
-    mock_mlflow_init.assert_called_once_with(
-        ANY,
-        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,
-    )
-
-    mock_minio_init.assert_called_once_with(
-        ANY,
+    mock_storage_repo.assert_called_once_with(
         endpoint_url=minio_config['endpoint_url'],
         access_key=minio_config['access_key'],
         secret_key=minio_config['secret_key'],
@@ -92,444 +112,196 @@ def test___init__(
         retry_mode=minio_config['retry_mode'],
         connect_timeout=minio_config['connect_timeout'],
         read_timeout=minio_config['read_timeout'],
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
     )
 
-    mock_training_init.assert_called_once_with(
-        ANY, logger=logger, notification_handler=notification_handler
-    )
+    mock_training_init.assert_called_once()
+    assert mock_training_init.call_args[1]['model_repository'] is mock_model_repo.return_value
+    assert mock_training_init.call_args[1]['storage_repository'] is mock_storage_repo.return_value
+    assert mock_training_init.call_args[1]['logger'] is mock_logger
+    assert mock_training_init.call_args[1]['notification_handler'] is mock_notification_handler
+
+    assert hasattr(activities, 'model_repository')
+    assert hasattr(activities, 'storage_repository')
 
 
-@mark.asyncio
-@patch('model_manager.activities.activities.ExperimentTracking', return_value=MagicMock())
-@patch('model_manager.activities.activities.MLFlow', return_value=MagicMock())
-async def test_shutdown(_mock_mlflow_init, mock_experiment_tracking_init):
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    await activities.shutdown()
-    mock_experiment_tracking_init.close.assert_called_once()
-
-
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___with_engine(
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.ExperimentTracking.close')
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_shutdown(
+    mock_storage_repo,
+    mock_model_repo,
     mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
+    mock_et_close,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
 ):
-    """Test __del__ calls parent destructor when engine attribute exists."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
+    """Test Activities.shutdown() calls ExperimentTracking.close()."""
+    from model_manager.activities.activities import Activities
 
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    # Add engine attribute to simulate Postgres initialization
-    activities.engine = MagicMock()
+    asyncio.run(activities.shutdown())
 
-    # Create a mock __del__ that will be detected by hasattr
-    mock_parent_del = MagicMock()
-
-    # Patch both the class and the instance to ensure super().__del__ exists and is callable
-    with patch.object(ExperimentTracking, '__del__', mock_parent_del, create=True):
-        # Trigger __del__
-        activities.__del__()
-
-        # Verify parent __del__ was called
-        mock_parent_del.assert_called_once()
+    mock_et_close.assert_called_once_with(activities)
 
 
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___without_engine(
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_del_without_engine(
+    mock_storage_repo,
+    mock_model_repo,
     mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
 ):
-    """Test __del__ does not call parent destructor when engine attribute is missing."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
+    """Test __del__ when engine attribute does not exist."""
+    from model_manager.activities.activities import Activities
 
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    # Ensure engine attribute does NOT exist
     if hasattr(activities, 'engine'):
         delattr(activities, 'engine')
 
-    # Mock super().__del__ to track if it's called
-    with patch.object(ExperimentTracking, '__del__', MagicMock()) as mock_parent_del:
-        # Trigger __del__
+    activities.__del__()
+
+
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_del_with_engine_no_super_del(
+    mock_storage_repo,
+    mock_model_repo,
+    mock_training_init,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test __del__ when engine exists but super has no __del__."""
+    from model_manager.activities.activities import Activities
+
+    activities = Activities(
+        postgres_config=postgres_config,
+        mlflow_config=mlflow_config,
+        minio_config=minio_config,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    activities.engine = MagicMock()
+
+    with patch('builtins.super') as mock_super:
+        mock_super_instance = MagicMock()
+        del mock_super_instance.__del__
+        mock_super.return_value = mock_super_instance
+
         activities.__del__()
 
-        # Verify parent __del__ was NOT called
-        mock_parent_del.assert_not_called()
 
-
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___handles_exception_gracefully(
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_del_with_engine_and_super_del(
+    mock_storage_repo,
+    mock_model_repo,
     mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
 ):
-    """Test __del__ handles exceptions from parent destructor gracefully."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
+    """Test __del__ when engine exists and super has __del__."""
+    from model_manager.activities.activities import Activities
 
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    # Add engine attribute
     activities.engine = MagicMock()
 
-    # Mock super().__del__ to raise an exception
-    mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed'))
+    mock_super_del = MagicMock()
 
-    with patch.object(ExperimentTracking, '__del__', mock_parent_del):
-        # Trigger __del__ - should not raise exception
-        try:
-            activities.__del__()
-            # Test passes if no exception is raised
-        except Exception as e:
-            # Test fails if exception propagates
-            raise AssertionError(f'__del__ should not raise exception, but raised: {e}') from e
+    class MockSuper:
+        def __del__(self):
+            mock_super_del()
 
-
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___when_parent_has_no_del(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """Test __del__ handles case when parent class has no __del__ method."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    # Add engine attribute
-    activities.engine = MagicMock()
-
-    # Remove __del__ from parent to simulate it not existing
-    with patch.object(ExperimentTracking, '__del__', create=False):
-        # Trigger __del__ - should not raise exception
-        try:
-            activities.__del__()
-            # Test passes if no exception is raised
-        except Exception as e:
-            # Test fails if exception propagates
-            raise AssertionError(
-                f'__del__ should handle missing parent __del__, but raised: {e}'
-            ) from e
-
-
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___calls_super_successfully(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """Test __del__ successfully calls super().__del__() when it exists - covers line 118."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    # Mock all parent __init__ methods to return None
-    mock_experiment_tracking_init.return_value = None
-    mock_mlflow_init.return_value = None
-    mock_minio_init.return_value = None
-    mock_training_init.return_value = None
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    # Add engine attribute to simulate Postgres initialization
-    activities.engine = MagicMock()
-
-    # Track if super().__del__() was actually called
-    super_del_called = []
-
-    def mock_super_del(self):
-        """Mock parent __del__ that tracks when it's called."""
-        super_del_called.append(True)
-
-    # Patch ExperimentTracking.__del__ to exist and be callable
-    with patch.object(ExperimentTracking, '__del__', mock_super_del, create=True):
-        # Trigger __del__ - this should execute line 118: super().__del__()
+    with patch('builtins.super', return_value=MockSuper()):
         activities.__del__()
 
-        # Verify that super().__del__() was actually called (line 118 executed)
-        assert len(super_del_called) == 1, 'super().__del__() should have been called once'
+        mock_super_del.assert_called_once()
 
 
-@patch('model_manager.activities.activities.ExperimentTracking.__init__')
-@patch('model_manager.activities.activities.MLFlow.__init__')
-@patch('model_manager.activities.activities.MinIO.__init__')
-@patch('model_manager.activities.activities.Training.__init__')
-def test___del___when_super_has_no_del_method(
+@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
+@patch('model_manager.activities.activities.Training.__init__', return_value=None)
+@patch('model_manager.activities.activities.ModelRepository')
+@patch('model_manager.activities.activities.StorageRepository')
+def test_activities_del_with_engine_exception_caught(
+    mock_storage_repo,
+    mock_model_repo,
     mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
+    mock_et_init,
+    postgres_config,
+    mlflow_config,
+    minio_config,
+    mock_logger,
+    mock_notification_handler,
 ):
-    """Test __del__ handles case when hasattr(super(), '__del__') returns False - covers line 118 false branch."""
-    postgres_config = {
-        'host': 'localhost',
-        'port': 5432,
-        'user': 'postgres',
-        'password': 'postgres',
-        'dbname': 'postgres',
-        'min_connections': 1,
-        'max_connections': 10,
-    }
-
-    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
-
-    minio_config = {
-        'endpoint_url': 'http://localhost:9000',
-        'access_key': 'minioadmin',
-        'secret_key': 'minioadmin',
-        'region': 'us-east-1',
-        'use_ssl': False,
-        'max_retry_attempts': 3,
-        'retry_mode': 'adaptive',
-        'connect_timeout': 10,
-        'read_timeout': 60,
-    }
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    # Mock all __init__ methods to return None
-    mock_experiment_tracking_init.return_value = None
-    mock_mlflow_init.return_value = None
-    mock_minio_init.return_value = None
-    mock_training_init.return_value = None
+    """Test __del__ catches exceptions when super().__del__() raises."""
+    from model_manager.activities.activities import Activities
 
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
         minio_config=minio_config,
-        logger=logger,
-        notification_handler=notification_handler,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    # Add engine attribute to pass the first hasattr check (line 115)
     activities.engine = MagicMock()
 
-    # Create a mock class without __del__ method to simulate super() not having __del__
-    class MockSuperWithoutDel:
-        """Mock class that explicitly does not have __del__ method."""
+    class MockSuperWithError:
+        def __del__(self):
+            raise RuntimeError('Test error')
 
-        pass
+    # Suppress the PytestUnraisableExceptionWarning for this specific test
+    import warnings
 
-    # Patch super() to return an instance that doesn't have __del__
-    mock_super_instance = MockSuperWithoutDel()
+    warnings.filterwarnings('ignore', category=pytest.PytestUnraisableExceptionWarning)
 
-    with patch('builtins.super', return_value=mock_super_instance):
-        # Trigger __del__ - should handle the case when hasattr(super(), '__del__') is False
-        try:
-            activities.__del__()
-            # Test passes - the false branch of line 118 was executed without error
-        except Exception as e:
-            # Test fails if exception propagates
-            raise AssertionError(
-                f'__del__ should handle super() without __del__ method, but raised: {e}'
-            ) from e
+    with patch('builtins.super', return_value=MockSuperWithError()):
+        activities.__del__()

tests/activities/test_minio.py

@@ -1,328 +0,0 @@
-from io import BytesIO
-from unittest.mock import MagicMock, patch
-
-from pytest import fixture, mark, raises
-
-from model_manager.activities.minio import MinIO
-
-
-@patch('model_manager.activities.minio.boto3.client')
-def test___init__(mock_boto3_client):
-    """Test MinIO initialization with correct configuration."""
-    mock_client = MagicMock()
-    mock_boto3_client.return_value = mock_client
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    minio = MinIO(
-        endpoint_url='http://localhost:9000',
-        access_key='minioadmin',
-        secret_key='minioadmin',
-        region='us-east-1',
-        use_ssl=False,
-        max_retry_attempts=3,
-        retry_mode='adaptive',
-        connect_timeout=10,
-        read_timeout=60,
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    assert minio.endpoint_url == 'http://localhost:9000'
-    assert minio.access_key == 'minioadmin'
-    assert minio.secret_key == 'minioadmin'
-    assert minio.region == 'us-east-1'
-    assert minio.use_ssl is False
-    assert minio.max_retry_attempts == 3
-    assert minio.retry_mode == 'adaptive'
-    assert minio.connect_timeout == 10
-    assert minio.read_timeout == 60
-
-    # Verify boto3 client was created with correct parameters
-    mock_boto3_client.assert_called_once()
-    call_kwargs = mock_boto3_client.call_args[1]
-    assert call_kwargs['endpoint_url'] == 'http://localhost:9000'
-    assert call_kwargs['aws_access_key_id'] == 'minioadmin'
-    assert call_kwargs['aws_secret_access_key'] == 'minioadmin'
-    assert call_kwargs['use_ssl'] is False
-
-
-@patch('model_manager.activities.minio.boto3.client')
-def test___init___failure(mock_boto3_client):
-    """Test MinIO initialization failure handling."""
-    mock_boto3_client.side_effect = Exception('Connection failed')
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    with raises(ConnectionError, match='Failed to initialize MinIO client'):
-        MinIO(
-            endpoint_url='http://localhost:9000',
-            access_key='minioadmin',
-            secret_key='minioadmin',
-            region='us-east-1',
-            use_ssl=False,
-            max_retry_attempts=3,
-            retry_mode='adaptive',
-            connect_timeout=10,
-            read_timeout=60,
-            logger=logger,
-            notification_handler=notification_handler,
-        )
-
-
-@fixture
-@patch('model_manager.activities.minio.boto3.client')
-def minio(mock_boto3_client):
-    """Fixture to create a MinIO instance for testing."""
-    mock_client = MagicMock()
-    mock_boto3_client.return_value = mock_client
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-
-    minio_instance = MinIO(
-        endpoint_url='http://localhost:9000',
-        access_key='minioadmin',
-        secret_key='minioadmin',
-        region='us-east-1',
-        use_ssl=False,
-        max_retry_attempts=3,
-        retry_mode='adaptive',
-        connect_timeout=10,
-        read_timeout=60,
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    minio_instance.send_notification = MagicMock()
-    minio_instance.minio_client = mock_client
-
-    return minio_instance
-
-
-metadata = {
-    'metadata': {
-        'workflow_name': 'test_workflow',
-        'model_name': 'test_model',
-        'model_id': 'test_model_id',
-    }
-}
-
-
-@mark.asyncio
-async def test_fetch_file_from_minio_success(minio):
-    """Test successful file fetch from MinIO."""
-    # Arrange
-    test_content = b'test file content'
-    mock_response = {'Body': MagicMock()}
-    mock_response['Body'].__enter__ = MagicMock(
-        return_value=MagicMock(read=MagicMock(return_value=test_content))
-    )
-    mock_response['Body'].__exit__ = MagicMock(return_value=None)
-
-    minio.minio_client.get_object.return_value = mock_response
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.txt',
-    }
-
-    # Act
-    result = await minio.fetch_file_from_minio(input_data)
-
-    # Assert
-    assert isinstance(result, BytesIO)
-    result.seek(0)
-    assert result.read() == test_content
-
-    minio.minio_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.txt')
-
-
-@mark.asyncio
-async def test_fetch_file_from_minio_file_not_found(minio):
-    """Test file fetch when file doesn't exist."""
-    # Arrange
-    minio.minio_client.get_object.side_effect = Exception(
-        'NoSuchKey: The specified key does not exist'
-    )
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'nonexistent.txt',
-    }
-
-    # Act & Assert
-    with raises(OSError, match='Error fetching file from MinIO'):
-        await minio.fetch_file_from_minio(input_data)
-
-    # Verify notification was sent
-    minio.send_notification.assert_called_once()
-    call_kwargs = minio.send_notification.call_args[1]
-    assert call_kwargs['notification_id'] == 'FETCH_FILE_FROM_MINIO_ERROR'
-    assert call_kwargs['block'] == 'fetch_file_from_minio'
-
-
-@mark.asyncio
-async def test_fetch_file_from_minio_network_error(minio):
-    """Test file fetch with network error."""
-    # Arrange
-    minio.minio_client.get_object.side_effect = Exception('Network timeout')
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.txt',
-    }
-
-    # Act & Assert
-    with raises(OSError, match='Error fetching file from MinIO'):
-        await minio.fetch_file_from_minio(input_data)
-
-    minio.send_notification.assert_called_once()
-
-
-@mark.asyncio
-async def test_delete_file_from_minio_success(minio):
-    """Test successful file deletion from MinIO."""
-    # Arrange
-    minio.minio_client.delete_object.return_value = None
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.txt',
-    }
-
-    # Act
-    result = await minio.delete_file_from_minio(input_data)
-
-    # Assert
-    assert result is None
-
-    minio.minio_client.delete_object.assert_called_once_with(
-        Bucket='test-bucket', Key='test-file.txt'
-    )
-
-
-@mark.asyncio
-async def test_delete_file_from_minio_idempotent(minio):
-    """Test that delete is idempotent (no error if file doesn't exist)."""
-    # Arrange
-    # MinIO delete_object is idempotent - no error if file doesn't exist
-    minio.minio_client.delete_object.return_value = None
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'nonexistent.txt',
-    }
-
-    # Act
-    result = await minio.delete_file_from_minio(input_data)
-
-    # Assert
-    assert result is None
-    minio.minio_client.delete_object.assert_called_once()
-
-
-@mark.asyncio
-async def test_delete_file_from_minio_access_denied(minio):
-    """Test file deletion with access denied error."""
-    # Arrange
-    minio.minio_client.delete_object.side_effect = Exception('AccessDenied: Access Denied')
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.txt',
-    }
-
-    # Act & Assert
-    with raises(OSError, match='Error deleting file from MinIO'):
-        await minio.delete_file_from_minio(input_data)
-
-    # Verify notification was sent
-    minio.send_notification.assert_called_once()
-    call_kwargs = minio.send_notification.call_args[1]
-    assert call_kwargs['notification_id'] == 'DELETE_FILE_FROM_MINIO_ERROR'
-    assert call_kwargs['block'] == 'delete_file_from_minio'
-
-
-@mark.asyncio
-async def test_delete_file_from_minio_network_error(minio):
-    """Test file deletion with network error."""
-    # Arrange
-    minio.minio_client.delete_object.side_effect = Exception('Connection timeout')
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.txt',
-    }
-
-    # Act & Assert
-    with raises(OSError, match='Error deleting file from MinIO'):
-        await minio.delete_file_from_minio(input_data)
-
-    minio.send_notification.assert_called_once()
-
-
-@mark.asyncio
-async def test_fetch_file_from_minio_large_file(minio):
-    """Test fetching a large file from MinIO."""
-    # Arrange
-    # Simulate a 10MB file
-    large_content = b'x' * (10 * 1024 * 1024)
-    mock_response = {'Body': MagicMock()}
-    mock_response['Body'].__enter__ = MagicMock(
-        return_value=MagicMock(read=MagicMock(return_value=large_content))
-    )
-    mock_response['Body'].__exit__ = MagicMock(return_value=None)
-
-    minio.minio_client.get_object.return_value = mock_response
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'large-file.bin',
-    }
-
-    # Act
-    result = await minio.fetch_file_from_minio(input_data)
-
-    # Assert
-    assert isinstance(result, BytesIO)
-    result.seek(0)
-    assert len(result.read()) == 10 * 1024 * 1024
-
-
-@mark.asyncio
-async def test_fetch_file_from_minio_empty_file(minio):
-    """Test fetching an empty file from MinIO."""
-    # Arrange
-    empty_content = b''
-    mock_response = {'Body': MagicMock()}
-    mock_response['Body'].__enter__ = MagicMock(
-        return_value=MagicMock(read=MagicMock(return_value=empty_content))
-    )
-    mock_response['Body'].__exit__ = MagicMock(return_value=None)
-
-    minio.minio_client.get_object.return_value = mock_response
-
-    input_data = {
-        'metadata': metadata['metadata'],
-        'bucket_name': 'test-bucket',
-        'file_name': 'empty-file.txt',
-    }
-
-    # Act
-    result = await minio.fetch_file_from_minio(input_data)
-
-    # Assert
-    assert isinstance(result, BytesIO)
-    result.seek(0)
-    assert result.read() == b''

tests/activities/test_mlflow.py

@@ -1,447 +0,0 @@
-from unittest.mock import ANY, MagicMock, patch
-
-import pytest
-from pytest import fixture, mark
-from sientia_do.notifications.models import NotificationLevel
-
-from model_manager.activities.mlflow import MLFlow
-
-
-@patch('model_manager.activities.mlflow.MLFlowRepository')
-def test___init__(mock_mlflow_repository):
-    mlflow = MLFlow(
-        mlflow_host='http://localhost',
-        mlflow_port=5000,
-        mlflow_username='admin',
-        mlflow_password='admin',
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    assert mlflow.mlflow_host == 'http://localhost'
-    assert mlflow.mlflow_port == 5000
-    assert mlflow.mlflow_username == 'admin'
-    assert mlflow.mlflow_password == 'admin'
-
-    mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
-
-
-@fixture
-@patch('model_manager.activities.mlflow.MLFlowRepository')
-def mlflow(mock_mlflow_repository):
-    mlflow = MLFlow(
-        mlflow_host='http://localhost:5000',
-        mlflow_port=5000,
-        mlflow_username='admin',
-        mlflow_password='admin',
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    mlflow.send_notification = MagicMock()
-
-    return mlflow
-
-
-metadata = {
-    'metadata': {
-        'model_id': 'test_model',
-        'model_name': 'test_model',
-        'workflow_name': 'test_workflow',
-        'schema_name': 'test_schedule',
-    },
-}
-
-
-@mark.asyncio
-async def test_save_model_success(mlflow):
-    """Test save_model successfully saves model and artifacts to MLflow."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    # Mock train result
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'test_experiment'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-    train_result.run_name = None  # Will be set by get_next_run_name
-
-    # Mock repository methods
-    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
-    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
-    mlflow.model_monitoring_repository.save_run.return_value = None
-
-    input_data = {
-        **metadata,
-        'train_result': train_result,
-    }
-
-    # Call the method
-    response = await mlflow.save_model(input_data)
-
-    # Verify repository methods were called
-    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
-    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
-    mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
-
-    # Verify response - now returns TrainModelResult directly
-    assert response == train_result
-    assert response.run_name == 'test_experiment-1'
-    assert train_result.run_name == 'test_experiment-1'
-
-
-@mark.asyncio
-async def test_save_model_get_next_run_name_error(mlflow):
-    """Test save_model handles error during get_next_run_name."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'test_experiment'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-
-    # Mock error in get_next_run_name
-    mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
-        'MLflow connection error'
-    )
-
-    input_data = {
-        **metadata,
-        'train_result': train_result,
-    }
-
-    # Call the method - should raise exception
-    with pytest.raises(Exception, match='MLflow connection error'):
-        await mlflow.save_model(input_data)
-
-    # Verify notification was sent
-    mlflow.send_notification.assert_called_once_with(
-        metadata=metadata['metadata'],
-        notification_id='SAVE_MODEL_ERROR',
-        message=ANY,
-        block='save_model',
-        level=NotificationLevel.ERROR,
-        attachment_content=ANY,
-    )
-
-
-@mark.asyncio
-async def test_save_model_generate_artifacts_error(mlflow):
-    """Test save_model handles error during generate_artifacts."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'test_experiment'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-
-    # Mock successful get_next_run_name but error in generate_artifacts
-    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
-    mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
-        'Reports directory does not exist'
-    )
-
-    input_data = {
-        **metadata,
-        'train_result': train_result,
-    }
-
-    # Call the method - should raise exception
-    with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
-        await mlflow.save_model(input_data)
-
-    # Verify notification was sent
-    mlflow.send_notification.assert_called_once_with(
-        metadata=metadata['metadata'],
-        notification_id='SAVE_MODEL_ERROR',
-        message=ANY,
-        block='save_model',
-        level=NotificationLevel.ERROR,
-        attachment_content=ANY,
-    )
-
-
-@mark.asyncio
-async def test_save_model_save_run_error(mlflow):
-    """Test save_model handles error during save_run."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'test_experiment'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-
-    # Mock successful get_next_run_name and generate_artifacts but error in save_run
-    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
-    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
-    mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
-        'One or more metrics (MSE, R2, MAE) are None'
-    )
-
-    input_data = {
-        **metadata,
-        'train_result': train_result,
-    }
-
-    # Call the method - should raise exception
-    with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
-        await mlflow.save_model(input_data)
-
-    # Verify notification was sent
-    mlflow.send_notification.assert_called_once_with(
-        metadata=metadata['metadata'],
-        notification_id='SAVE_MODEL_ERROR',
-        message=ANY,
-        block='save_model',
-        level=NotificationLevel.ERROR,
-        attachment_content=ANY,
-    )
-
-
-@mark.asyncio
-async def test_save_model_missing_metadata(mlflow):
-    """Test save_model handles missing metadata gracefully."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'test_experiment'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-
-    # Mock repository methods
-    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
-    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
-    mlflow.model_monitoring_repository.save_run.return_value = None
-
-    # Input data without metadata
-    input_data = {
-        'train_result': train_result,
-    }
-
-    # Call the method
-    response = await mlflow.save_model(input_data)
-
-    # Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
-    assert response == train_result
-    assert response.run_name == 'test_experiment-1'
-
-
-@mark.asyncio
-async def test_save_model_complete_flow(mlflow):
-    """Test save_model complete flow with all steps."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    params.experiment_name = 'production_model'
-
-    train_result = MagicMock(spec=TrainModelResult)
-    train_result.params = params
-    train_result.run_name = None
-    train_result.run_dir = None
-    train_result.report_path = None
-
-    # Mock complete flow
-    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'
-
-    # After generate_artifacts, paths should be set
-    updated_result = MagicMock(spec=TrainModelResult)
-    updated_result.params = params
-    updated_result.run_name = 'production_model-5'
-    updated_result.run_dir = '/reports/production_model-5_20231010'
-    updated_result.report_path = '/reports/production_model-5_20231010/report.html'
-    updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
-    updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'
-
-    mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
-    mlflow.model_monitoring_repository.save_run.return_value = None
-
-    input_data = {
-        **metadata,
-        'train_result': train_result,
-    }
-
-    # Call the method
-    response = await mlflow.save_model(input_data)
-
-    # Verify complete flow
-    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
-    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
-    mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
-
-    # Verify response - returns TrainModelResult directly
-    assert response == updated_result
-    assert response.run_name == 'production_model-5'
-    assert response.run_dir is not None
-    assert response.report_path is not None
-
-
-# ============================================================================
-# Tests for cleanup_run_directory
-# ============================================================================
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_success(mlflow):
-    """Test successful cleanup of run directory."""
-    input_data = {
-        **metadata,
-        'run_dir': 'test_run_dir',
-    }
-
-    # Patch os and shutil inside the activity method
-    with (
-        patch('os.path.exists', return_value=True) as mock_exists,
-        patch('shutil.rmtree') as mock_rmtree,
-    ):
-        # Call the method
-        await mlflow.cleanup_run_directory(input_data)
-
-        # Verify directory existence was checked
-        mock_exists.assert_called_once_with('test_run_dir')
-
-        # Verify shutil.rmtree was called
-        mock_rmtree.assert_called_once_with('test_run_dir')
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_already_deleted(mlflow):
-    """Test cleanup when directory is already deleted (idempotent)."""
-    input_data = {
-        **metadata,
-        'run_dir': 'already_deleted_dir',
-    }
-
-    # Patch os and shutil inside the activity method
-    with (
-        patch('os.path.exists', return_value=False) as mock_exists,
-        patch('shutil.rmtree') as mock_rmtree,
-    ):
-        # Call the method - should not raise error
-        await mlflow.cleanup_run_directory(input_data)
-
-        # Verify directory existence was checked
-        mock_exists.assert_called_once_with('already_deleted_dir')
-
-        # Verify shutil.rmtree was NOT called
-        mock_rmtree.assert_not_called()
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_no_run_dir(mlflow):
-    """Test cleanup when no run_dir is provided."""
-    input_data = {
-        **metadata,
-        # No run_dir key
-    }
-
-    # Call the method - should not raise error
-    await mlflow.cleanup_run_directory(input_data)
-
-    # Should complete without errors
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_empty_run_dir(mlflow):
-    """Test cleanup when run_dir is empty string."""
-    input_data = {
-        **metadata,
-        'run_dir': '',
-    }
-
-    # Call the method - should not raise error
-    await mlflow.cleanup_run_directory(input_data)
-
-    # Should complete without errors
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_none_run_dir(mlflow):
-    """Test cleanup when run_dir is None."""
-    input_data = {
-        **metadata,
-        'run_dir': None,
-    }
-
-    # Call the method - should not raise error
-    await mlflow.cleanup_run_directory(input_data)
-
-    # Should complete without errors
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_error(mlflow):
-    """Test cleanup handles errors correctly."""
-    input_data = {
-        **metadata,
-        'run_dir': 'error_dir',
-    }
-
-    # Patch os and shutil inside the activity method
-    with (
-        patch('os.path.exists', return_value=True),
-        patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
-    ):
-        # Call the method - should raise exception
-        with pytest.raises(PermissionError, match='Permission denied'):
-            await mlflow.cleanup_run_directory(input_data)
-
-        # Verify notification was sent
-        mlflow.send_notification.assert_called_once_with(
-            metadata=metadata['metadata'],
-            notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
-            message=ANY,
-            block='cleanup_run_directory',
-            level=NotificationLevel.ERROR,
-            attachment_content=ANY,
-        )
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_missing_metadata(mlflow):
-    """Test cleanup handles missing metadata gracefully."""
-    input_data = {
-        'run_dir': 'no_metadata_dir',
-    }
-
-    # Patch os and shutil inside the activity method
-    with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
-        # Call the method
-        await mlflow.cleanup_run_directory(input_data)
-
-        # Verify directory was deleted
-        mock_rmtree.assert_called_once_with('no_metadata_dir')
-
-
-@mark.asyncio
-async def test_cleanup_run_directory_oserror(mlflow):
-    """Test cleanup handles OSError correctly."""
-    input_data = {
-        **metadata,
-        'run_dir': 'os_error_dir',
-    }
-
-    # Patch os and shutil inside the activity method
-    with (
-        patch('os.path.exists', return_value=True),
-        patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
-    ):
-        # Call the method - should raise exception
-        with pytest.raises(OSError, match='Directory not empty'):
-            await mlflow.cleanup_run_directory(input_data)
-
-        # Verify notification was sent
-        mlflow.send_notification.assert_called_once()
-        call_args = mlflow.send_notification.call_args[1]
-        assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
-        assert call_args['level'] == NotificationLevel.ERROR
-        assert 'Directory not empty' in call_args['message']

tests/activities/test_training.py

@@ -1,499 +1,497 @@
-"""Unit tests for Training activity."""
+"""Unit tests for Training class with 100% coverage."""
 
+import asyncio
 from io import BytesIO
 from unittest.mock import MagicMock, patch
 
 import pytest
-from pytest import mark
-
-from model_manager.activities.training import Training
-from model_manager.utils.models.train_model_params import TrainModelParams
-from model_manager.utils.models.train_model_result import TrainModelResult
 
 
-@mark.asyncio
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_notification_handler():
+    """Create a mock notification handler."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_model_repository():
+    """Create a mock model repository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_storage_repository():
+    """Create a mock storage repository."""
+    return MagicMock()
+
+
+@pytest.fixture
+def mock_train_params():
+    """Create a mock TrainModelParams."""
+    params = MagicMock()
+    params.experiment_run_id = 1
+    params.target_variable = 'target'
+    params.experiment_name = 'test_experiment'
+    params.bucket_name = 'test-bucket'
+    params.file_name = 'test-file.csv'
+    params.validate_business_rules = MagicMock()
+    return params
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
 @patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_success(mock_training_repository_class):
-    """Test successful model training."""
-    # Create mock repository instance
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
+def test_training_init(
+    mock_training_repo,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test Training initialization."""
+    from model_manager.activities.training import Training
 
-    # Create mock train result
-    mock_train_result = MagicMock(spec=TrainModelResult)
-    mock_train_result.mse_val = 0.5
-    mock_train_result.mae_val = 0.3
-    mock_train_result.r2_val = 0.95
-
-    mock_final_result = MagicMock(spec=TrainModelResult)
-    mock_final_result.mse_val = 0.5
-    mock_final_result.mae_val = 0.3
-    mock_final_result.r2_val = 0.95
-
-    # Setup repository mocks
-    mock_repository.train.return_value = mock_train_result
-    mock_repository.after_train_calculation.return_value = mock_final_result
-
-    # Create Training instance
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    # Mock inherited methods
-    training.info = MagicMock()
-
-    # Test data
-    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
-    train_params = TrainModelParams(
-        experiment_run_id=123,
-        target_variable='price',
-        variable_columns=['feature1', 'feature2'],
-        train_size=80,
-        shuffle=True,
-        use_scaler=True,
-        include_ar=False,
-        bucket_name='test-bucket',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        lag_train=1,
-        lag_val=1,
-        rem_static_win=False,
-        low_lim={'feature1': 0.0, 'feature2': 0.0},
-        upp_lim={'feature1': 100.0, 'feature2': 100.0},
-        window=10,
-        experiment_name='test_experiment',
-        removed_intervals=[],
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
+    mock_base_init.assert_called_once_with(
+        mock_logger, mock_notification_handler, set_error_counter=True
+    )
+    mock_training_repo.assert_called_once_with(mock_logger)
+    assert training.model_repository is mock_model_repository
+    assert training.storage_repository is mock_storage_repository
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+@patch('model_manager.activities.training.TrainModelParams')
+def test_validate_train_params_success(
+    mock_train_params_class,
+    mock_training_repo,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+    mock_train_params,
+):
+    """Test validate_train_params with valid parameters."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.info = MagicMock()
+    mock_train_params_class.from_dict.return_value = mock_train_params
+
     input_data = {
         'metadata': {'workflow_id': 'test-123'},
-        'uploaded_file': uploaded_file,
-        'train_params': train_params,
+        'experiment_run_id': 1,
+        'target_variable': 'target',
     }
 
-    # Execute
-    result = await training.train_model(input_data)
+    result = asyncio.run(training.validate_train_params(input_data))
 
-    # Assertions - now returns TrainModelResult directly
-    assert result == mock_final_result
-    assert result.mse_val == 0.5
-    assert result.mae_val == 0.3
-    assert result.r2_val == 0.95
-
-    # Verify repository calls
-    mock_repository.train.assert_called_once()
-    mock_repository.after_train_calculation.assert_called_once()
+    assert result is mock_train_params
+    mock_train_params_class.from_dict.assert_called_once_with(input_data)
+    mock_train_params.validate_business_rules.assert_called_once()
+    training.info.assert_called_once()
 
 
-@mark.asyncio
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
 @patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_invalid_file_type(mock_training_repository_class):
-    """Test training with invalid file type."""
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
+@patch('model_manager.activities.training.TrainModelParams')
+def test_validate_train_params_value_error(
+    mock_train_params_class,
+    mock_training_repo,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test validate_train_params with ValueError."""
+    from model_manager.activities.training import Training
 
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    # Invalid file type (string instead of BytesIO)
-    train_params = TrainModelParams(
-        experiment_run_id=123,
-        target_variable='price',
-        variable_columns=['feature1'],
-        train_size=80,
-        shuffle=True,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        lag_train=1,
-        lag_val=1,
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        experiment_name='test_experiment',
-        removed_intervals=[],
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
     )
 
-    input_data = {
-        'metadata': {},
-        'uploaded_file': 'not_a_bytesio',
-        'train_params': train_params,
-    }
-
-    # Should raise ValueError
-    with pytest.raises(ValueError, match='uploaded_file must be BytesIO'):
-        await training.train_model(input_data)
-
-    # Verify notification was sent (via BaseActivity)
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_training_error(mock_training_repository_class):
-    """Test training failure during model training."""
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
-
-    # Setup repository to raise error
-    mock_repository.train.side_effect = ValueError('Training data is empty')
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    uploaded_file = BytesIO(b'test,data\n')
-    train_params = TrainModelParams(
-        experiment_run_id=456,
-        target_variable='price',
-        variable_columns=['feature1'],
-        train_size=80,
-        shuffle=True,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        lag_train=1,
-        lag_val=1,
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        experiment_name='test_experiment',
-        removed_intervals=[],
-    )
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-456'},
-        'uploaded_file': uploaded_file,
-        'train_params': train_params,
-    }
-
-    # Should raise ValueError
-    with pytest.raises(ValueError, match='Training data is empty'):
-        await training.train_model(input_data)
-
-    # Verify notification was sent (via BaseActivity)
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_sends_notification_on_error(mock_training_repository_class):
-    """Test that notification is sent when training fails."""
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
-
-    mock_repository.train.side_effect = Exception('Database connection failed')
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    uploaded_file = BytesIO(b'test,data\n1,2')
-    train_params = TrainModelParams(
-        experiment_run_id=789,
-        target_variable='price',
-        variable_columns=['feature1'],
-        train_size=80,
-        shuffle=True,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        lag_train=1,
-        lag_val=1,
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        experiment_name='test_experiment',
-        removed_intervals=[],
-    )
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-789', 'experiment_run_id': 789},
-        'uploaded_file': uploaded_file,
-        'train_params': train_params,
-    }
-
-    # Should raise Exception
-    with pytest.raises(Exception, match='Database connection failed'):
-        await training.train_model(input_data)
-
-    # Verify notification was sent (via BaseActivity)
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_after_calculation_error(mock_training_repository_class):
-    """Test training failure during post-training calculations."""
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
-
-    # Train succeeds but after_calculation fails
-    mock_train_result = MagicMock(spec=TrainModelResult)
-    mock_repository.train.return_value = mock_train_result
-    mock_repository.after_train_calculation.side_effect = Exception('Metric calculation failed')
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
-    train_params = TrainModelParams(
-        experiment_run_id=999,
-        target_variable='price',
-        variable_columns=['feature1'],
-        train_size=80,
-        shuffle=True,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        lag_train=1,
-        lag_val=1,
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        experiment_name='test_experiment',
-        removed_intervals=[],
-    )
-
-    input_data = {
-        'metadata': {},
-        'uploaded_file': uploaded_file,
-        'train_params': train_params,
-    }
-
-    # Should raise Exception
-    with pytest.raises(Exception, match='Metric calculation failed'):
-        await training.train_model(input_data)
-
-    # Verify notification was sent (via BaseActivity)
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.training.TrainingRepository')
-async def test_train_model_invalid_train_params_type(mock_training_repository_class):
-    """Test training with invalid train_params type (dict instead of TrainModelParams)."""
-    mock_repository = MagicMock()
-    mock_training_repository_class.return_value = mock_repository
-
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    uploaded_file = BytesIO(b'test,data\n1,2\n3,4')
-
-    # Invalid train_params type (dict instead of TrainModelParams object)
-    input_data = {
-        'metadata': {'workflow_id': 'test-invalid'},
-        'uploaded_file': uploaded_file,
-        'train_params': {
-            'experiment_run_id': 123,
-            'target_variable': 'price',
-            'variable_columns': ['feature1'],
-        },  # This is a dict, not TrainModelParams
-    }
-
-    # Should raise ValueError
-    with pytest.raises(ValueError, match='train_params must be TrainModelParams.*dict'):
-        await training.train_model(input_data)
-
-
-# ============================================================================
-# Tests for validate_train_params
-# ============================================================================
-
-
-@mark.asyncio
-async def test_validate_train_params_success():
-    """Test successful validation of training parameters."""
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    training.info = MagicMock()
+    training.send_notification = MagicMock()
+    mock_train_params_class.from_dict.side_effect = ValueError('Invalid parameter')
 
     input_data = {
         'metadata': {'workflow_id': 'test-123'},
-        'experiment_run_id': 456,
-        'target_variable': 'price',
-        'variable_columns': ['feature1', 'feature2', 'price'],  # target_variable must be in list
-        'train_size': 80,
-        'shuffle': True,
-        'use_scaler': True,
-        'include_ar': False,
-        'bucket_name': 'test-bucket',
-        'file_name': 'test.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'lag_train': 1,
-        'lag_val': 1,
-        'rem_static_win': False,
-        'low_lim': {'feature1': 0.0, 'feature2': 0.0, 'price': 0.0},
-        'upp_lim': {'feature1': 100.0, 'feature2': 100.0, 'price': 1000.0},
-        'window': 10,
-        'experiment_name': 'test_experiment',
-        'removed_intervals': [],
+        'experiment_run_id': 1,
     }
 
-    result = await training.validate_train_params(input_data)
-
-    assert isinstance(result, TrainModelParams)
-    assert result.experiment_run_id == 456
-    assert result.target_variable == 'price'
-    assert result.variable_columns == ['feature1', 'feature2', 'price']
-    assert result.train_size == 80
-    assert result.experiment_name == 'test_experiment'
-    assert training.info.call_count == 2
-
-
-@mark.asyncio
-async def test_validate_train_params_missing_required_field():
-    """Test validation fails when required field is missing."""
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    training.info = MagicMock()
-    training.error = MagicMock()
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-123'},
-        'experiment_run_id': 456,
-        'variable_columns': ['feature1'],
-        'train_size': 80,
-        'shuffle': True,
-        'use_scaler': False,
-        'include_ar': False,
-        'bucket_name': 'test',
-        'file_name': 'test.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'lag_train': 1,
-        'lag_val': 1,
-        'rem_static_win': False,
-        'low_lim': {'feature1': 0.0},
-        'upp_lim': {'feature1': 100.0},
-        'window': 10,
-        'experiment_name': 'test',
-        'removed_intervals': [],
-    }
-
-    with pytest.raises(ValueError, match='target_variable'):
-        await training.validate_train_params(input_data)
-
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-async def test_validate_train_params_invalid_type():
-    """Test validation fails when field has invalid type."""
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    training.info = MagicMock()
-    training.error = MagicMock()
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-invalid'},
-        'experiment_run_id': 456,
-        'target_variable': 'price',
-        'variable_columns': ['feature1'],
-        'train_size': 'invalid',
-        'shuffle': True,
-        'use_scaler': False,
-        'include_ar': False,
-        'bucket_name': 'test',
-        'file_name': 'test.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'lag_train': 1,
-        'lag_val': 1,
-        'rem_static_win': False,
-        'low_lim': {'feature1': 0.0},
-        'upp_lim': {'feature1': 100.0},
-        'window': 10,
-        'experiment_name': 'test',
-        'removed_intervals': [],
-    }
-
-    with pytest.raises((ValueError, TypeError)):
-        await training.validate_train_params(input_data)
-
-    notification_handler.send_notification.assert_called_once()
-
-
-@mark.asyncio
-async def test_validate_train_params_empty_input():
-    """Test validation fails with empty input."""
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
-
-    training.info = MagicMock()
-    training.error = MagicMock()
-
-    input_data = {'metadata': {}}
-
     with pytest.raises(ValueError):
-        await training.validate_train_params(input_data)
+        asyncio.run(training.validate_train_params(input_data))
 
-    notification_handler.send_notification.assert_called_once()
+    training.send_notification.assert_called_once()
 
 
-@mark.asyncio
-async def test_validate_train_params_without_metadata():
-    """Test validation works even without metadata key."""
-    logger = MagicMock()
-    notification_handler = MagicMock()
-    training = Training(logger=logger, notification_handler=notification_handler)
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+@patch('model_manager.activities.training.TrainModelParams')
+def test_validate_train_params_type_error(
+    mock_train_params_class,
+    mock_training_repo,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test validate_train_params with TypeError."""
+    from model_manager.activities.training import Training
 
-    training.info = MagicMock()
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.send_notification = MagicMock()
+    mock_train_params_class.from_dict.side_effect = TypeError('Type mismatch')
 
     input_data = {
-        'experiment_run_id': 789,
-        'target_variable': 'temperature',
-        'variable_columns': ['sensor1', 'temperature'],  # target_variable must be in list
-        'train_size': 75,
-        'shuffle': False,
-        'use_scaler': True,
-        'include_ar': True,
-        'bucket_name': 'sensors',
-        'file_name': 'data.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'lag_train': 2,
-        'lag_val': 2,
-        'rem_static_win': True,
-        'low_lim': {'sensor1': -50.0, 'temperature': -50.0},
-        'upp_lim': {'sensor1': 150.0, 'temperature': 150.0},
-        'window': 20,
-        'experiment_name': 'sensor_experiment',
-        'removed_intervals': [],
+        'metadata': {'workflow_id': 'test-123'},
+        'experiment_run_id': 1,
     }
 
-    result = await training.validate_train_params(input_data)
+    with pytest.raises(TypeError):
+        asyncio.run(training.validate_train_params(input_data))
 
-    assert isinstance(result, TrainModelParams)
-    assert result.experiment_run_id == 789
-    assert result.target_variable == 'temperature'
-    assert result.experiment_name == 'sensor_experiment'
+    training.send_notification.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+@patch('model_manager.activities.training.TrainModelParams')
+def test_validate_train_params_key_error(
+    mock_train_params_class,
+    mock_training_repo,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test validate_train_params with KeyError."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.send_notification = MagicMock()
+    mock_train_params_class.from_dict.side_effect = KeyError('missing_key')
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'experiment_run_id': 1,
+    }
+
+    with pytest.raises(KeyError):
+        asyncio.run(training.validate_train_params(input_data))
+
+    training.send_notification.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_train_model_success_with_params_object(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+    mock_train_params,
+):
+    """Test train_model with TrainModelParams object."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    mock_file = BytesIO(b'test data')
+    mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
+
+    mock_train_result = MagicMock()
+    mock_train_result.run_name = 'run_001'
+    mock_train_result.run_dir = '/tmp/run_001'  # noqa: S108
+
+    training.training_repository.train.return_value = mock_train_result
+    training.training_repository.after_train_calculation.return_value = mock_train_result
+    mock_model_repository.save_model.return_value = mock_train_result
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'train_params': mock_train_params,
+    }
+
+    result = asyncio.run(training.train_model(input_data))
+
+    assert result == {'run_name': 'run_001', 'run_dir': '/tmp/run_001'}  # noqa: S108
+    mock_storage_repository.fetch_file.assert_called_once_with('test-bucket', 'test-file.csv')
+    training.training_repository.train.assert_called_once_with(mock_file, mock_train_params)
+    training.training_repository.after_train_calculation.assert_called_once_with(
+        mock_train_params, mock_train_result
+    )
+    mock_model_repository.save_model.assert_called_once_with(mock_train_result)
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+@patch('model_manager.activities.training.TrainModelParams')
+def test_train_model_success_with_params_dict(
+    mock_train_params_class,
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+    mock_train_params,
+):
+    """Test train_model with dict parameters."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    mock_train_params_class.from_dict.return_value = mock_train_params
+    mock_file = BytesIO(b'test data')
+    mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
+
+    mock_train_result = MagicMock()
+    mock_train_result.run_name = 'run_002'
+    mock_train_result.run_dir = '/tmp/run_002'  # noqa: S108
+
+    training.training_repository.train.return_value = mock_train_result
+    training.training_repository.after_train_calculation.return_value = mock_train_result
+    mock_model_repository.save_model.return_value = mock_train_result
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'train_params': {'experiment_run_id': 1, 'target_variable': 'target'},
+    }
+
+    result = asyncio.run(training.train_model(input_data))
+
+    assert result == {'run_name': 'run_002', 'run_dir': '/tmp/run_002'}  # noqa: S108
+    mock_train_params_class.from_dict.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_train_model_training_fails(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+    mock_train_params,
+):
+    """Test train_model when training fails."""
+    from model_manager.activities.training import Training
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.send_notification = MagicMock()
+    mock_file = BytesIO(b'test data')
+    mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
+    training.training_repository.train.side_effect = RuntimeError('Training failed')
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'train_params': mock_train_params,
+    }
+
+    with pytest.raises(ModelTrainingError) as exc_info:
+        asyncio.run(training.train_model(input_data))
+
+    assert exc_info.value.model_trained is False
+    assert exc_info.value.model_saved is False
+    training.send_notification.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_train_model_save_fails(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+    mock_train_params,
+):
+    """Test train_model when model saving fails."""
+    from model_manager.activities.training import Training
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.send_notification = MagicMock()
+    mock_file = BytesIO(b'test data')
+    mock_storage_repository.fetch_file.return_value.__enter__.return_value = mock_file
+
+    mock_train_result = MagicMock()
+    training.training_repository.train.return_value = mock_train_result
+    training.training_repository.after_train_calculation.return_value = mock_train_result
+    mock_model_repository.save_model.side_effect = RuntimeError('Save failed')
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'train_params': mock_train_params,
+    }
+
+    with pytest.raises(ModelTrainingError) as exc_info:
+        asyncio.run(training.train_model(input_data))
+
+    assert exc_info.value.model_trained is True
+    assert exc_info.value.model_saved is False
+    training.send_notification.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_cleanup_resources_success(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test cleanup_resources successfully."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'run_dir': '/tmp/run_001',  # noqa: S108
+        'bucket_name': 'test-bucket',
+        'file_name': 'test-file.csv',
+    }
+
+    asyncio.run(training.cleanup_resources(input_data))
+
+    mock_model_repository.cleanup_run_directory.assert_called_once_with('/tmp/run_001')  # noqa: S108
+    mock_storage_repository.delete_file.assert_called_once_with('test-bucket', 'test-file.csv')
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_cleanup_resources_cleanup_fails(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test cleanup_resources when cleanup fails."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    training.send_notification = MagicMock()
+    mock_model_repository.cleanup_run_directory.side_effect = RuntimeError('Cleanup failed')
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+        'run_dir': '/tmp/run_001',  # noqa: S108
+        'bucket_name': 'test-bucket',
+        'file_name': 'test-file.csv',
+    }
+
+    with pytest.raises(RuntimeError):
+        asyncio.run(training.cleanup_resources(input_data))
+
+    training.send_notification.assert_called_once()
+
+
+@patch('model_manager.activities.training.BaseActivity.__init__', return_value=None)
+@patch('model_manager.activities.training.TrainingRepository')
+def test_cleanup_resources_with_empty_values(
+    mock_training_repo_class,
+    mock_base_init,
+    mock_model_repository,
+    mock_storage_repository,
+    mock_logger,
+    mock_notification_handler,
+):
+    """Test cleanup_resources with empty values."""
+    from model_manager.activities.training import Training
+
+    training = Training(
+        model_repository=mock_model_repository,
+        storage_repository=mock_storage_repository,
+        logger=mock_logger,
+        notification_handler=mock_notification_handler,
+    )
+
+    input_data = {
+        'metadata': {'workflow_id': 'test-123'},
+    }
+
+    asyncio.run(training.cleanup_resources(input_data))
+
+    mock_model_repository.cleanup_run_directory.assert_called_once_with('')
+    mock_storage_repository.delete_file.assert_called_once_with('', '')

tests/sientia/test_model_serving.py

@@ -16,9 +16,8 @@ def test_init_with_all_credentials(mock_set_tracking_uri):
     tracking_uri = 'http://mlflow.example.com'
     username = 'test_user'
     password = 'test_pass'
-    logger = MagicMock()
 
-    ModelServing(tracking_uri=tracking_uri, username=username, password=password, logger=logger)
+    ModelServing(tracking_uri=tracking_uri, username=username, password=password)
 
     mock_set_tracking_uri.assert_called_once_with(tracking_uri)
     import os

tests/test_metrics.py

@@ -0,0 +1,209 @@
+"""Unit tests for model_manager.metrics module.
+
+This module tests the Prometheus metrics configuration used for
+monitoring and observability in the Sientia DataOps Model Manager.
+"""
+
+
+def test_app_up_metric_exists():
+    """Test that APP_UP metric is properly defined."""
+    from model_manager.metrics import APP_UP
+
+    assert APP_UP is not None
+    assert APP_UP._name == 'app_up'
+    assert (
+        APP_UP._documentation == 'Indicates if the application is running (1) or shutting down (0)'
+    )
+
+
+def test_app_up_metric_has_pod_id_label():
+    """Test that APP_UP metric has pod_id label."""
+    from model_manager.metrics import APP_UP
+
+    assert 'pod_id' in APP_UP._labelnames
+
+
+def test_app_up_metric_is_gauge():
+    """Test that APP_UP is a Gauge metric."""
+    from prometheus_client import Gauge
+
+    from model_manager.metrics import APP_UP
+
+    assert isinstance(APP_UP, Gauge)
+
+
+def test_app_up_metric_can_be_set_to_one():
+    """Test that APP_UP metric can be set to 1 (running)."""
+    from model_manager.metrics import APP_UP
+
+    # Set metric to 1 for a specific pod
+    APP_UP.labels(pod_id='test-pod-1').set(1)
+
+    # Verify the metric value
+    metric_value = APP_UP.labels(pod_id='test-pod-1')._value._value
+    assert metric_value == 1
+
+
+def test_app_up_metric_can_be_set_to_zero():
+    """Test that APP_UP metric can be set to 0 (shutting down)."""
+    from model_manager.metrics import APP_UP
+
+    # Set metric to 0 for a specific pod
+    APP_UP.labels(pod_id='test-pod-2').set(0)
+
+    # Verify the metric value
+    metric_value = APP_UP.labels(pod_id='test-pod-2')._value._value
+    assert metric_value == 0
+
+
+def test_app_up_metric_multiple_pods():
+    """Test that APP_UP metric can track multiple pods independently."""
+    from model_manager.metrics import APP_UP
+
+    # Set different values for different pods
+    APP_UP.labels(pod_id='pod-1').set(1)
+    APP_UP.labels(pod_id='pod-2').set(0)
+    APP_UP.labels(pod_id='pod-3').set(1)
+
+    # Verify each pod has correct value
+    assert APP_UP.labels(pod_id='pod-1')._value._value == 1
+    assert APP_UP.labels(pod_id='pod-2')._value._value == 0
+    assert APP_UP.labels(pod_id='pod-3')._value._value == 1
+
+
+def test_app_up_metric_default_value():
+    """Test that APP_UP metric starts with no value set."""
+    # Create a new label that hasn't been used yet
+    import uuid
+
+    from model_manager.metrics import APP_UP
+
+    unique_pod = f'test-pod-{uuid.uuid4()}'
+
+    # The metric should exist but not have a value until set
+    metric = APP_UP.labels(pod_id=unique_pod)
+    assert metric is not None
+
+
+def test_metrics_module_imports():
+    """Test that metrics module can be imported successfully."""
+    import model_manager.metrics
+
+    assert hasattr(model_manager.metrics, 'APP_UP')
+    assert hasattr(model_manager.metrics, 'Gauge')
+
+
+def test_metrics_module_docstring():
+    """Test that metrics module has proper documentation."""
+    import model_manager.metrics
+
+    assert model_manager.metrics.__doc__ is not None
+    assert 'Prometheus' in model_manager.metrics.__doc__
+    assert 'metrics' in model_manager.metrics.__doc__
+
+
+def test_app_up_metric_can_increment():
+    """Test that APP_UP metric value can be incremented."""
+    from model_manager.metrics import APP_UP
+
+    pod_id = 'test-pod-increment'
+    APP_UP.labels(pod_id=pod_id).set(0)
+
+    # Increment the metric
+    APP_UP.labels(pod_id=pod_id).inc()
+
+    metric_value = APP_UP.labels(pod_id=pod_id)._value._value
+    assert metric_value == 1
+
+
+def test_app_up_metric_can_decrement():
+    """Test that APP_UP metric value can be decremented."""
+    from model_manager.metrics import APP_UP
+
+    pod_id = 'test-pod-decrement'
+    APP_UP.labels(pod_id=pod_id).set(1)
+
+    # Decrement the metric
+    APP_UP.labels(pod_id=pod_id).dec()
+
+    metric_value = APP_UP.labels(pod_id=pod_id)._value._value
+    assert metric_value == 0
+
+
+def test_app_up_metric_set_to_timestamp():
+    """Test that APP_UP metric can be set to current timestamp."""
+    import time
+
+    from model_manager.metrics import APP_UP
+
+    pod_id = 'test-pod-timestamp'
+    current_time = time.time()
+
+    # Set to timestamp
+    APP_UP.labels(pod_id=pod_id).set_to_current_time()
+
+    metric_value = APP_UP.labels(pod_id=pod_id)._value._value
+
+    # Should be close to current time
+    assert abs(metric_value - current_time) < 2  # Within 2 seconds
+
+
+def test_app_up_metric_label_validation():
+    """Test that APP_UP metric validates label names."""
+    from model_manager.metrics import APP_UP
+
+    # Should work with valid label
+    APP_UP.labels(pod_id='valid-pod-name').set(1)
+
+    # Should work with empty string (though not recommended)
+    APP_UP.labels(pod_id='').set(1)
+
+    # Should work with special characters
+    APP_UP.labels(pod_id='pod-123_test.example').set(1)
+
+
+def test_module_exports():
+    """Test that metrics module exports expected symbols."""
+    import model_manager.metrics as metrics_module
+
+    # Check that module has the expected exports
+    module_contents = dir(metrics_module)
+
+    assert 'APP_UP' in module_contents
+    assert 'Gauge' in module_contents
+
+
+def test_app_up_metric_thread_safety():
+    """Test that APP_UP metric is thread-safe."""
+    import threading
+
+    from model_manager.metrics import APP_UP
+
+    pod_id = 'test-pod-threading'
+    APP_UP.labels(pod_id=pod_id).set(0)
+
+    def increment_metric():
+        for _ in range(100):
+            APP_UP.labels(pod_id=pod_id).inc()
+
+    # Create multiple threads that increment the metric
+    threads = [threading.Thread(target=increment_metric) for _ in range(5)]
+
+    for thread in threads:
+        thread.start()
+
+    for thread in threads:
+        thread.join()
+
+    # Should have incremented 500 times total
+    metric_value = APP_UP.labels(pod_id=pod_id)._value._value
+    assert metric_value == 500
+
+
+def test_prometheus_client_gauge_import():
+    """Test that Gauge is properly imported from prometheus_client."""
+    from prometheus_client import Gauge as PrometheusGauge
+
+    from model_manager.metrics import Gauge
+
+    assert Gauge is PrometheusGauge

tests/utils/models/test_train_model_params.py

@@ -1,497 +1,371 @@
-"""Unit tests for TrainModelParams class."""
+"""Unit tests for TrainModelParams with 100% coverage."""
 
 import pytest
 
-from model_manager.utils.models.train_model_params import TrainModelParams
-
 
 @pytest.fixture
-def valid_params_dict():
-    """Create valid parameters dictionary for testing."""
+def valid_train_params_dict():
+    """Create a valid dictionary for TrainModelParams."""
     return {
-        'variable_columns': ['var1', 'var2', 'var3'],
+        'variable_columns': ['var1', 'var2'],
         'lag_train': 5,
         'lag_val': 3,
         'target_variable': 'target',
         'rem_static_win': True,
-        'low_lim': {'var1': 0.0, 'var2': 0.0, 'var3': 0.0},
-        'upp_lim': {'var1': 100.0, 'var2': 100.0, 'var3': 100.0},
+        'low_lim': {'var1': 0.0, 'var2': 1.0},
+        'upp_lim': {'var1': 10.0, 'var2': 20.0},
         'window': 10,
         'use_scaler': True,
         'include_ar': False,
         'bucket_name': 'test-bucket',
         'file_name': 'test-file.csv',
-        'line_separator': '\n',
+        'line_separator': ',',
         'decimal_separator': '.',
         'train_size': 80,
         'shuffle': True,
-        'experiment_run_id': 123,
-        'experiment_name': 'Test experiment name',
+        'experiment_run_id': 1,
+        'experiment_name': 'test_experiment',
         'removed_intervals': [],
     }
 
 
-def test_train_model_params_creation_with_valid_params(valid_params_dict):
-    """Test creating TrainModelParams with all valid parameters."""
-    params = TrainModelParams(**valid_params_dict)
+def test_train_model_params_from_dict_success(valid_train_params_dict):
+    """Test TrainModelParams.from_dict with valid data."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    assert params.variable_columns == ['var1', 'var2', 'var3']
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    assert params.variable_columns == ['var1', 'var2']
     assert params.lag_train == 5
     assert params.lag_val == 3
     assert params.target_variable == 'target'
     assert params.rem_static_win is True
-    assert params.low_lim == {'var1': 0.0, 'var2': 0.0, 'var3': 0.0}
-    assert params.upp_lim == {'var1': 100.0, 'var2': 100.0, 'var3': 100.0}
+    assert params.low_lim == {'var1': 0.0, 'var2': 1.0}
+    assert params.upp_lim == {'var1': 10.0, 'var2': 20.0}
     assert params.window == 10
     assert params.use_scaler is True
     assert params.include_ar is False
     assert params.bucket_name == 'test-bucket'
     assert params.file_name == 'test-file.csv'
-    assert params.line_separator == '\n'
+    assert params.line_separator == ','
     assert params.decimal_separator == '.'
     assert params.train_size == 80
     assert params.shuffle is True
-    assert params.experiment_run_id == 123
-    assert params.experiment_name == 'Test experiment name'
+    assert params.experiment_run_id == 1
+    assert params.experiment_name == 'test_experiment'
     assert params.removed_intervals == []
 
 
-def test_train_model_params_from_dict_creation(valid_params_dict):
-    """Test creating TrainModelParams using from_dict method."""
-    params = TrainModelParams.from_dict(valid_params_dict)
+def test_train_model_params_check_none_raises_value_error():
+    """Test _check_none raises ValueError when value is None."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    assert params.variable_columns == ['var1', 'var2', 'var3']
-    assert params.lag_train == 5
-    assert params.experiment_run_id == 123
+    with pytest.raises(ValueError, match='test_field is required and cannot be None'):
+        TrainModelParams._check_none(None, str, 'test_field')
 
 
-def test_train_model_params_variable_columns_none_raises_error(valid_params_dict):
-    """Test that None variable_columns raises ValueError."""
-    valid_params_dict['variable_columns'] = None
+def test_train_model_params_check_none_raises_type_error():
+    """Test _check_none raises TypeError when type is incorrect."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    with pytest.raises(TypeError, match='test_field must be of type str, but got int'):
+        TrainModelParams._check_none(123, str, 'test_field')
+
+
+def test_train_model_params_check_none_success():
+    """Test _check_none returns value when valid."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    result = TrainModelParams._check_none('test_value', str, 'test_field')
+    assert result == 'test_value'
+
+
+def test_train_model_params_check_type_raises_type_error():
+    """Test _check_type raises TypeError when type is incorrect."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    with pytest.raises(TypeError, match='test_field must be of type int, but got str'):
+        TrainModelParams._check_type('not_an_int', int, 'test_field')
+
+
+def test_train_model_params_check_type_success():
+    """Test _check_type returns value when valid."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    result = TrainModelParams._check_type(42, int, 'test_field')
+    assert result == 42
+
+
+def test_train_model_params_check_type_with_none():
+    """Test _check_type allows None value."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    result = TrainModelParams._check_type(None, str, 'test_field')
+    assert result is None
+
+
+def test_train_model_params_from_dict_missing_field(valid_train_params_dict):
+    """Test from_dict raises ValueError when required field is missing."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    del valid_train_params_dict['variable_columns']
 
     with pytest.raises(ValueError, match='variable_columns is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
+        TrainModelParams.from_dict(valid_train_params_dict)
 
 
-def test_train_model_params_variable_columns_wrong_type_raises_error(valid_params_dict):
-    """Test that wrong type for variable_columns raises TypeError."""
-    valid_params_dict['variable_columns'] = 'not a list'
+def test_train_model_params_from_dict_wrong_type(valid_train_params_dict):
+    """Test from_dict raises TypeError when field has wrong type."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    with pytest.raises(TypeError, match='variable_columns must be of type list'):
-        TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['lag_train'] = 'not_an_int'
+
+    with pytest.raises(TypeError, match='lag_train must be of type int, but got str'):
+        TrainModelParams.from_dict(valid_train_params_dict)
 
 
-def test_train_model_params_lag_train_none_raises_error(valid_params_dict):
-    """Test that None lag_train raises ValueError."""
-    valid_params_dict['lag_train'] = None
+def test_train_model_params_from_dict_with_none_removed_intervals(valid_train_params_dict):
+    """Test from_dict allows None for removed_intervals."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    with pytest.raises(ValueError, match='lag_train is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['removed_intervals'] = None
 
-
-def test_train_model_params_lag_train_wrong_type_raises_error(valid_params_dict):
-    """Test that wrong type for lag_train raises TypeError."""
-    valid_params_dict['lag_train'] = '5'
-
-    with pytest.raises(TypeError, match='lag_train must be of type int'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_target_variable_none_raises_error(valid_params_dict):
-    """Test that None target_variable raises ValueError."""
-    valid_params_dict['target_variable'] = None
-
-    with pytest.raises(ValueError, match='target_variable is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_target_variable_wrong_type_raises_error(valid_params_dict):
-    """Test that wrong type for target_variable raises TypeError."""
-    valid_params_dict['target_variable'] = 123
-
-    with pytest.raises(TypeError, match='target_variable must be of type str'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_boolean_fields(valid_params_dict):
-    """Test boolean fields validation."""
-    # Test rem_static_win
-    valid_params_dict['rem_static_win'] = None
-    with pytest.raises(ValueError, match='rem_static_win is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-    valid_params_dict['rem_static_win'] = 'true'
-    with pytest.raises(TypeError, match='rem_static_win must be of type bool'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_dict_fields(valid_params_dict):
-    """Test dict fields validation."""
-    # Test low_lim
-    valid_params_dict['low_lim'] = None
-    with pytest.raises(ValueError, match='low_lim is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-    valid_params_dict['low_lim'] = {'var1': 0.0}
-    valid_params_dict['upp_lim'] = 'not a dict'
-    with pytest.raises(TypeError, match='upp_lim must be of type dict'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_bucket_name_none_raises_error(valid_params_dict):
-    """Test that None bucket_name raises ValueError."""
-    valid_params_dict['bucket_name'] = None
-
-    with pytest.raises(ValueError, match='bucket_name is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_file_name_none_raises_error(valid_params_dict):
-    """Test that None file_name raises ValueError."""
-    valid_params_dict['file_name'] = None
-
-    with pytest.raises(ValueError, match='file_name is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_experiment_run_id_none_raises_error(valid_params_dict):
-    """Test that None experiment_run_id raises ValueError."""
-    valid_params_dict['experiment_run_id'] = None
-
-    with pytest.raises(ValueError, match='experiment_run_id is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_experiment_name_none_raises_error(valid_params_dict):
-    """Test that None experiment_name raises ValueError."""
-    valid_params_dict['experiment_name'] = None
-
-    with pytest.raises(ValueError, match='experiment_name is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-def test_train_model_params_removed_intervals_can_be_none(valid_params_dict):
-    """Test that removed_intervals can be None (uses _check_type not _check_none)."""
-    valid_params_dict['removed_intervals'] = None
-
-    params = TrainModelParams.from_dict(valid_params_dict)
+    params = TrainModelParams.from_dict(valid_train_params_dict)
     assert params.removed_intervals is None
 
 
-def test_train_model_params_removed_intervals_wrong_type_raises_error(valid_params_dict):
-    """Test that wrong type for removed_intervals raises TypeError."""
-    valid_params_dict['removed_intervals'] = 'not a list'
+def test_validate_business_rules_success(valid_train_params_dict):
+    """Test validate_business_rules with valid parameters."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    with pytest.raises(TypeError, match='removed_intervals must be of type list'):
-        TrainModelParams.from_dict(valid_params_dict)
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+    params.validate_business_rules()  # Should not raise
 
 
-def test_train_model_params_removed_intervals_with_values(valid_params_dict):
-    """Test removed_intervals with actual interval values."""
-    valid_params_dict['removed_intervals'] = [
-        ('2023-01-01', '2023-01-10'),
-        ('2023-02-01', '2023-02-05'),
-    ]
+def test_validate_business_rules_train_size_too_low(valid_train_params_dict):
+    """Test validate_business_rules raises error when train_size < 10."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
-    assert len(params.removed_intervals) == 2
-    assert params.removed_intervals[0] == ('2023-01-01', '2023-01-10')
+    valid_train_params_dict['train_size'] = 5
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-
-def test_train_model_params_all_fields_count():
-    """Test that TrainModelParams has exactly 19 required fields."""
-    import inspect
-
-    sig = inspect.signature(TrainModelParams.__init__)
-    # Subtract 1 for 'self'
-    param_count = len(sig.parameters) - 1
-    assert param_count == 19
-
-
-def test_train_model_params_with_minimal_valid_data():
-    """Test creating params with minimal valid data."""
-    params = TrainModelParams(
-        variable_columns=['x'],
-        lag_train=1,
-        lag_val=1,
-        target_variable='y',
-        rem_static_win=False,
-        low_lim={},
-        upp_lim={},
-        window=1,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='bucket',
-        file_name='file.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        train_size=50,
-        shuffle=False,
-        experiment_run_id=1,
-        experiment_name='name',
-        removed_intervals=[],
-    )
-
-    assert params.variable_columns == ['x']
-    assert params.lag_train == 1
-    assert params.experiment_run_id == 1
-
-
-def test_train_model_params_check_none_method():
-    """Test _check_none method behavior."""
-    params_dict = {
-        'variable_columns': ['var1'],
-        'lag_train': 5,
-        'lag_val': 3,
-        'target_variable': 'target',
-        'rem_static_win': True,
-        'low_lim': {},
-        'upp_lim': {},
-        'window': 10,
-        'use_scaler': True,
-        'include_ar': False,
-        'bucket_name': 'bucket',
-        'file_name': 'file.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'train_size': 80,
-        'shuffle': True,
-        'experiment_run_id': 123,
-        'experiment_name': 'exp',
-        'removed_intervals': [],
-    }
-
-    params = TrainModelParams(**params_dict)
-
-    # Test that _check_none is a private method
-    assert hasattr(params, '_check_none')
-    assert callable(params._check_none)
-
-
-def test_train_model_params_check_type_method():
-    """Test _check_type method behavior."""
-    params_dict = {
-        'variable_columns': ['var1'],
-        'lag_train': 5,
-        'lag_val': 3,
-        'target_variable': 'target',
-        'rem_static_win': True,
-        'low_lim': {},
-        'upp_lim': {},
-        'window': 10,
-        'use_scaler': True,
-        'include_ar': False,
-        'bucket_name': 'bucket',
-        'file_name': 'file.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'train_size': 80,
-        'shuffle': True,
-        'experiment_run_id': 123,
-        'experiment_name': 'exp',
-        'removed_intervals': [],
-    }
-
-    params = TrainModelParams(**params_dict)
-
-    # Test that _check_type is a private method
-    assert hasattr(params, '_check_type')
-    assert callable(params._check_type)
-
-
-# ============================================================================
-# Tests for validate_business_rules method
-# ============================================================================
-
-
-def test_validate_business_rules_success(valid_params_dict):
-    """Test that valid params pass business rules validation."""
-    # Ensure target_variable is in variable_columns
-    valid_params_dict['target_variable'] = 'var1'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    # Should not raise any exception
-    params.validate_business_rules()
-
-
-def test_validate_business_rules_train_size_too_low(valid_params_dict):
-    """Test that train_size < 1 raises ValueError."""
-    valid_params_dict['train_size'] = 0
-    valid_params_dict['target_variable'] = 'var1'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='train_size must be between 1 and 99'):
+    with pytest.raises(ValueError, match='train_size must be between 10 and 100, got 5'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_train_size_too_high(valid_params_dict):
-    """Test that train_size > 99 raises ValueError."""
-    valid_params_dict['train_size'] = 100
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_train_size_too_high(valid_train_params_dict):
+    """Test validate_business_rules raises error when train_size > 100."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['train_size'] = 101
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='train_size must be between 1 and 99'):
+    with pytest.raises(ValueError, match='train_size must be between 10 and 100, got 101'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_empty_variable_columns(valid_params_dict):
-    """Test that empty variable_columns raises ValueError."""
-    valid_params_dict['variable_columns'] = []
+def test_validate_business_rules_empty_variable_columns(valid_train_params_dict):
+    """Test validate_business_rules raises error when variable_columns is empty."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['variable_columns'] = []
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='variable_columns cannot be empty'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_lag_train_zero(valid_params_dict):
-    """Test that lag_train = 0 raises ValueError."""
-    valid_params_dict['lag_train'] = 0
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_negative_lag_train(valid_train_params_dict):
+    """Test validate_business_rules raises error when lag_train is negative."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['lag_train'] = -1
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='lag_train must be positive'):
+    with pytest.raises(ValueError, match='lag_train must be positive, got -1'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_lag_train_negative(valid_params_dict):
-    """Test that lag_train < 0 raises ValueError."""
-    valid_params_dict['lag_train'] = -1
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_negative_lag_val(valid_train_params_dict):
+    """Test validate_business_rules raises error when lag_val is negative."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['lag_val'] = -2
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='lag_train must be positive'):
+    with pytest.raises(ValueError, match='lag_val must be positive, got -2'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_lag_val_zero(valid_params_dict):
-    """Test that lag_val = 0 raises ValueError."""
-    valid_params_dict['lag_val'] = 0
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_negative_window(valid_train_params_dict):
+    """Test validate_business_rules raises error when window is negative."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['window'] = -5
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='lag_val must be positive'):
+    with pytest.raises(ValueError, match='window must be positive, got -5'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_window_zero(valid_params_dict):
-    """Test that window = 0 raises ValueError."""
-    valid_params_dict['window'] = 0
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_mismatched_limit_keys(valid_train_params_dict):
+    """Test validate_business_rules raises error when low_lim and upp_lim keys don't match."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='window must be positive'):
-        params.validate_business_rules()
-
-
-def test_validate_business_rules_low_lim_upp_lim_keys_mismatch(valid_params_dict):
-    """Test that mismatched keys in low_lim and upp_lim raises ValueError."""
-    valid_params_dict['low_lim'] = {'var1': 0.0, 'var2': 0.0}
-    valid_params_dict['upp_lim'] = {'var1': 100.0, 'var3': 100.0}  # var3 instead of var2
-    valid_params_dict['target_variable'] = 'var1'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['low_lim'] = {'var1': 0.0}
+    valid_train_params_dict['upp_lim'] = {'var1': 10.0, 'var2': 20.0}
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='low_lim and upp_lim must have the same keys'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_low_lim_greater_than_upp_lim(valid_params_dict):
-    """Test that low_lim >= upp_lim raises ValueError."""
-    valid_params_dict['low_lim'] = {'var1': 100.0, 'var2': 0.0, 'var3': 0.0}
-    valid_params_dict['upp_lim'] = {'var1': 50.0, 'var2': 100.0, 'var3': 100.0}
-    valid_params_dict['target_variable'] = 'var2'
+def test_validate_business_rules_low_lim_greater_than_upp_lim(valid_train_params_dict):
+    """Test validate_business_rules raises error when low_lim >= upp_lim."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['low_lim'] = {'var1': 15.0, 'var2': 1.0}
+    valid_train_params_dict['upp_lim'] = {'var1': 10.0, 'var2': 20.0}
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='low_lim must be less than upp_lim for variable "var1"'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_low_lim_equal_to_upp_lim(valid_params_dict):
-    """Test that low_lim == upp_lim raises ValueError."""
-    valid_params_dict['low_lim'] = {'var1': 50.0, 'var2': 0.0, 'var3': 0.0}
-    valid_params_dict['upp_lim'] = {'var1': 50.0, 'var2': 100.0, 'var3': 100.0}
-    valid_params_dict['target_variable'] = 'var2'
+def test_validate_business_rules_low_lim_equal_to_upp_lim(valid_train_params_dict):
+    """Test validate_business_rules raises error when low_lim == upp_lim."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['low_lim'] = {'var1': 10.0, 'var2': 1.0}
+    valid_train_params_dict['upp_lim'] = {'var1': 10.0, 'var2': 20.0}
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='low_lim must be less than upp_lim for variable "var1"'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_target_not_in_variable_columns(valid_params_dict):
-    """Test that target_variable not in variable_columns raises ValueError."""
-    valid_params_dict['target_variable'] = 'nonexistent_var'
+def test_validate_business_rules_empty_bucket_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when bucket_name is empty."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(
-        ValueError, match='target_variable "nonexistent_var" must be in variable_columns'
-    ):
-        params.validate_business_rules()
-
-
-def test_validate_business_rules_empty_bucket_name(valid_params_dict):
-    """Test that empty bucket_name raises ValueError."""
-    valid_params_dict['bucket_name'] = '   '
-    valid_params_dict['target_variable'] = 'var1'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['bucket_name'] = ''
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='bucket_name cannot be empty or whitespace'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_empty_file_name(valid_params_dict):
-    """Test that empty file_name raises ValueError."""
-    valid_params_dict['file_name'] = ''
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_whitespace_bucket_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when bucket_name is whitespace."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['bucket_name'] = '   '
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    with pytest.raises(ValueError, match='bucket_name cannot be empty or whitespace'):
+        params.validate_business_rules()
+
+
+def test_validate_business_rules_empty_file_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when file_name is empty."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['file_name'] = ''
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='file_name cannot be empty or whitespace'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_empty_experiment_name(valid_params_dict):
-    """Test that empty experiment_name raises ValueError."""
-    valid_params_dict['experiment_name'] = '   '
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_whitespace_file_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when file_name is whitespace."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['file_name'] = '  \t  '
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    with pytest.raises(ValueError, match='file_name cannot be empty or whitespace'):
+        params.validate_business_rules()
+
+
+def test_validate_business_rules_empty_experiment_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when experiment_name is empty."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['experiment_name'] = ''
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
     with pytest.raises(ValueError, match='experiment_name cannot be empty or whitespace'):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_all_valid_edge_cases(valid_params_dict):
-    """Test that edge case valid values pass validation."""
-    valid_params_dict['train_size'] = 1  # Minimum valid
-    valid_params_dict['lag_train'] = 1  # Minimum valid
-    valid_params_dict['lag_val'] = 1  # Minimum valid
-    valid_params_dict['window'] = 1  # Minimum valid
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_whitespace_experiment_name(valid_train_params_dict):
+    """Test validate_business_rules raises error when experiment_name is whitespace."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['experiment_name'] = '  \n  '
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    # Should not raise any exception
-    params.validate_business_rules()
+    with pytest.raises(ValueError, match='experiment_name cannot be empty or whitespace'):
+        params.validate_business_rules()
 
 
-def test_validate_business_rules_train_size_99(valid_params_dict):
-    """Test that train_size = 99 (maximum valid) passes validation."""
-    valid_params_dict['train_size'] = 99
-    valid_params_dict['target_variable'] = 'var1'
+def test_validate_business_rules_train_size_boundary_10(valid_train_params_dict):
+    """Test validate_business_rules accepts train_size = 10 (lower boundary)."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
 
-    params = TrainModelParams.from_dict(valid_params_dict)
+    valid_train_params_dict['train_size'] = 10
+    params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    # Should not raise any exception
-    params.validate_business_rules()
+    params.validate_business_rules()  # Should not raise
+
+
+def test_validate_business_rules_train_size_boundary_100(valid_train_params_dict):
+    """Test validate_business_rules accepts train_size = 100 (upper boundary)."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['train_size'] = 100
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    params.validate_business_rules()  # Should not raise
+
+
+def test_validate_business_rules_zero_lag_train(valid_train_params_dict):
+    """Test validate_business_rules accepts lag_train = 0."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['lag_train'] = 0
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    params.validate_business_rules()  # Should not raise
+
+
+def test_validate_business_rules_zero_lag_val(valid_train_params_dict):
+    """Test validate_business_rules accepts lag_val = 0."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['lag_val'] = 0
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    params.validate_business_rules()  # Should not raise
+
+
+def test_validate_business_rules_zero_window(valid_train_params_dict):
+    """Test validate_business_rules accepts window = 0."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['window'] = 0
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    params.validate_business_rules()  # Should not raise
+
+
+def test_validate_business_rules_empty_limits(valid_train_params_dict):
+    """Test validate_business_rules accepts empty low_lim and upp_lim."""
+    from model_manager.utils.models.train_model_params import TrainModelParams
+
+    valid_train_params_dict['low_lim'] = {}
+    valid_train_params_dict['upp_lim'] = {}
+    params = TrainModelParams.from_dict(valid_train_params_dict)
+
+    params.validate_business_rules()  # Should not raise

tests/utils/models/test_train_model_result.py

@@ -175,11 +175,11 @@ def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
 
 
 def test_train_model_result_field_count():
-    """Test that TrainModelResult has exactly 17 fields."""
+    """Test that TrainModelResult has exactly 19 fields."""
     from dataclasses import fields
 
     result_fields = fields(TrainModelResult)
-    assert len(result_fields) == 17
+    assert len(result_fields) == 19
 
     field_names = {f.name for f in result_fields}
     expected_fields = {
@@ -195,6 +195,8 @@ def test_train_model_result_field_count():
         'mse_val',
         'mae_val',
         'r2_val',
+        'equation',
+        'equation_path',
         'run_name',
         'report_path',
         'train_data_path',

tests/utils/repository/test_storage_repository.py

@@ -0,0 +1,464 @@
+"""Unit tests for StorageRepository class."""
+
+from io import BytesIO
+from unittest.mock import Mock, patch
+
+import pytest
+from botocore.exceptions import ClientError
+
+
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger for testing."""
+    logger = Mock()
+    logger.info = Mock()
+    logger.error = Mock()
+    logger.warning = Mock()
+    return logger
+
+
+@pytest.fixture
+def storage_config():
+    """Create storage repository configuration."""
+    return {
+        'endpoint_url': 'http://localhost:9000',
+        'access_key': 'test_access_key',
+        'secret_key': 'test_secret_key',
+        'region': 'us-east-1',
+        'use_ssl': False,
+        'max_retry_attempts': 3,
+        'retry_mode': 'standard',
+        'connect_timeout': 30,
+        'read_timeout': 60,
+    }
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_initialization(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository initialization with correct parameters."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Verify attributes are set correctly
+    assert repo.endpoint_url == storage_config['endpoint_url']
+    assert repo.access_key == storage_config['access_key']
+    assert repo.secret_key == storage_config['secret_key']
+    assert repo.region == storage_config['region']
+    assert repo.use_ssl == storage_config['use_ssl']
+    assert repo.max_retry_attempts == storage_config['max_retry_attempts']
+    assert repo.retry_mode == storage_config['retry_mode']
+    assert repo.connect_timeout == storage_config['connect_timeout']
+    assert repo.read_timeout == storage_config['read_timeout']
+    assert repo.logger == mock_logger
+
+    # Verify boto3 client was created
+    mock_boto3.client.assert_called_once()
+    call_args = mock_boto3.client.call_args
+
+    assert call_args[0][0] == 's3'
+    assert call_args[1]['endpoint_url'] == storage_config['endpoint_url']
+    assert call_args[1]['aws_access_key_id'] == storage_config['access_key']
+    assert call_args[1]['aws_secret_access_key'] == storage_config['secret_key']
+    assert call_args[1]['use_ssl'] == storage_config['use_ssl']
+
+    # Verify logger was called
+    mock_logger.info.assert_called()
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_boto_config(mock_boto3, mock_logger, storage_config):
+    """Test that boto3 Config is created with correct retry settings."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    StorageRepository(logger=mock_logger, **storage_config)
+
+    # Verify Config was passed with correct settings
+    call_args = mock_boto3.client.call_args
+    boto_config = call_args[1]['config']
+
+    assert boto_config.region_name == storage_config['region']
+    assert boto_config.connect_timeout == storage_config['connect_timeout']
+    assert boto_config.read_timeout == storage_config['read_timeout']
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_success(mock_boto3, mock_logger, storage_config):
+    """Test successful file fetch from MinIO."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock response
+    file_content = b'test file content'
+    mock_body = Mock()
+    mock_body.read.return_value = file_content
+    mock_body.__enter__ = Mock(return_value=mock_body)
+    mock_body.__exit__ = Mock(return_value=False)
+
+    mock_response = {'Body': mock_body}
+
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.return_value = mock_response
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Fetch file
+    result = repo.fetch_file('test-bucket', 'test-file.csv')
+
+    # Verify result
+    assert isinstance(result, BytesIO)
+    assert result.getvalue() == file_content
+
+    # Verify get_object was called correctly
+    mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
+
+    # Verify logging
+    assert mock_logger.info.call_count >= 2  # Init + fetch
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_with_large_content(mock_boto3, mock_logger, storage_config):
+    """Test fetch file with large content."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock response with large content
+    large_content = b'x' * 1024 * 1024  # 1MB
+    mock_body = Mock()
+    mock_body.read.return_value = large_content
+    mock_body.__enter__ = Mock(return_value=mock_body)
+    mock_body.__exit__ = Mock(return_value=False)
+
+    mock_response = {'Body': mock_body}
+
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.return_value = mock_response
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Fetch file
+    result = repo.fetch_file('test-bucket', 'large-file.bin')
+
+    # Verify result
+    assert isinstance(result, BytesIO)
+    assert len(result.getvalue()) == 1024 * 1024
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_empty_content(mock_boto3, mock_logger, storage_config):
+    """Test fetch file with empty content."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock response with empty content
+    mock_body = Mock()
+    mock_body.read.return_value = b''
+    mock_body.__enter__ = Mock(return_value=mock_body)
+    mock_body.__exit__ = Mock(return_value=False)
+
+    mock_response = {'Body': mock_body}
+
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.return_value = mock_response
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Fetch file
+    result = repo.fetch_file('test-bucket', 'empty-file.txt')
+
+    # Verify result
+    assert isinstance(result, BytesIO)
+    assert result.getvalue() == b''
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_not_found(mock_boto3, mock_logger, storage_config):
+    """Test fetch file when object doesn't exist."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock to raise NoSuchKey error
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.side_effect = ClientError(
+        {'Error': {'Code': 'NoSuchKey', 'Message': 'The specified key does not exist.'}},
+        'GetObject',
+    )
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Attempt to fetch non-existent file
+    with pytest.raises(ClientError) as exc_info:
+        repo.fetch_file('test-bucket', 'non-existent.csv')
+
+    assert exc_info.value.response['Error']['Code'] == 'NoSuchKey'
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_access_denied(mock_boto3, mock_logger, storage_config):
+    """Test fetch file when access is denied."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock to raise AccessDenied error
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.side_effect = ClientError(
+        {'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'GetObject'
+    )
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Attempt to fetch file without permissions
+    with pytest.raises(ClientError) as exc_info:
+        repo.fetch_file('test-bucket', 'protected-file.csv')
+
+    assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_network_error(mock_boto3, mock_logger, storage_config):
+    """Test fetch file when network error occurs."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock to raise network error
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.side_effect = ConnectionError('Network unreachable')
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Attempt to fetch file with network error
+    with pytest.raises(ConnectionError):
+        repo.fetch_file('test-bucket', 'test-file.csv')
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_delete_file_success(mock_boto3, mock_logger, storage_config):
+    """Test successful file deletion from MinIO."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_s3_client.delete_object.return_value = {}
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Delete file
+    repo.delete_file('test-bucket', 'test-file.csv')
+
+    # Verify delete_object was called correctly
+    mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
+
+    # Verify logging
+    assert mock_logger.info.call_count >= 2  # Init + delete
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_delete_file_non_existent(mock_boto3, mock_logger, storage_config):
+    """Test delete file that doesn't exist (should succeed silently in S3)."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # S3/MinIO delete is idempotent - deleting non-existent file succeeds
+    mock_s3_client = Mock()
+    mock_s3_client.delete_object.return_value = {}
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Delete non-existent file (should succeed)
+    repo.delete_file('test-bucket', 'non-existent.csv')
+
+    mock_s3_client.delete_object.assert_called_once()
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_delete_file_access_denied(mock_boto3, mock_logger, storage_config):
+    """Test delete file when access is denied."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock to raise AccessDenied error
+    mock_s3_client = Mock()
+    mock_s3_client.delete_object.side_effect = ClientError(
+        {'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'DeleteObject'
+    )
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Attempt to delete file without permissions
+    with pytest.raises(ClientError) as exc_info:
+        repo.delete_file('test-bucket', 'protected-file.csv')
+
+    assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_delete_file_network_error(mock_boto3, mock_logger, storage_config):
+    """Test delete file when network error occurs."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    # Setup mock to raise network error
+    mock_s3_client = Mock()
+    mock_s3_client.delete_object.side_effect = ConnectionError('Network unreachable')
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Attempt to delete file with network error
+    with pytest.raises(ConnectionError):
+        repo.delete_file('test-bucket', 'test-file.csv')
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_with_ssl(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository initialization with SSL enabled."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    storage_config['use_ssl'] = True
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    assert repo.use_ssl is True
+
+    # Verify boto3 client was created with use_ssl=True
+    call_args = mock_boto3.client.call_args
+    assert call_args[1]['use_ssl'] is True
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_custom_timeouts(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository with custom timeout values."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    storage_config['connect_timeout'] = 10
+    storage_config['read_timeout'] = 120
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    assert repo.connect_timeout == 10
+    assert repo.read_timeout == 120
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_custom_retry_mode(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository with different retry modes."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    for retry_mode in ['standard', 'legacy', 'adaptive']:
+        storage_config['retry_mode'] = retry_mode
+        repo = StorageRepository(logger=mock_logger, **storage_config)
+        assert repo.retry_mode == retry_mode
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_custom_max_retries(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository with different max retry attempts."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    storage_config['max_retry_attempts'] = 5
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    assert repo.max_retry_attempts == 5
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_with_special_characters(mock_boto3, mock_logger, storage_config):
+    """Test fetch file with special characters in name."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    file_content = b'test content'
+    mock_body = Mock()
+    mock_body.read.return_value = file_content
+    mock_body.__enter__ = Mock(return_value=mock_body)
+    mock_body.__exit__ = Mock(return_value=False)
+
+    mock_response = {'Body': mock_body}
+
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.return_value = mock_response
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Fetch file with special characters
+    special_filename = 'test file (2023-01-01) #1.csv'
+    result = repo.fetch_file('test-bucket', special_filename)
+
+    assert isinstance(result, BytesIO)
+    mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key=special_filename)
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_delete_file_with_path_separators(mock_boto3, mock_logger, storage_config):
+    """Test delete file with path separators in object key."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_s3_client.delete_object.return_value = {}
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    # Delete file with path separators
+    file_path = 'data/2023/01/test-file.csv'
+    repo.delete_file('test-bucket', file_path)
+
+    mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key=file_path)
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_storage_repository_different_regions(mock_boto3, mock_logger, storage_config):
+    """Test StorageRepository with different AWS regions."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    mock_s3_client = Mock()
+    mock_boto3.client.return_value = mock_s3_client
+
+    regions = ['us-west-1', 'eu-central-1', 'ap-southeast-1']
+
+    for region in regions:
+        storage_config['region'] = region
+        repo = StorageRepository(logger=mock_logger, **storage_config)
+        assert repo.region == region
+
+
+@patch('model_manager.utils.repository.storage_repository.boto3')
+def test_fetch_file_logs_file_size(mock_boto3, mock_logger, storage_config):
+    """Test that fetch_file logs the file size."""
+    from model_manager.utils.repository.storage_repository import StorageRepository
+
+    file_content = b'x' * 12345
+    mock_body = Mock()
+    mock_body.read.return_value = file_content
+    mock_body.__enter__ = Mock(return_value=mock_body)
+    mock_body.__exit__ = Mock(return_value=False)
+
+    mock_response = {'Body': mock_body}
+
+    mock_s3_client = Mock()
+    mock_s3_client.get_object.return_value = mock_response
+    mock_boto3.client.return_value = mock_s3_client
+
+    repo = StorageRepository(logger=mock_logger, **storage_config)
+
+    repo.fetch_file('test-bucket', 'test-file.csv')
+
+    # Verify logging includes file size
+    log_calls = [str(call) for call in mock_logger.info.call_args_list]
+    assert any('12345 bytes' in str(call) for call in log_calls)

tests/utils/test_connectors_config.py

@@ -10,8 +10,7 @@ from model_manager.utils.connectors_config import (
 
 def test_build_mlflow_config_with_env_vars():
     # Arrange
-    environ['MLFLOW_HOST'] = 'http://test-host'
-    environ['MLFLOW_PORT'] = '8080'
+    environ['MLFLOW_URL'] = 'http://test-host:8080'
     environ['MLFLOW_USERNAME'] = 'test-user'
     environ['MLFLOW_PASSWORD'] = 'test-pass'
 
@@ -19,8 +18,7 @@ def test_build_mlflow_config_with_env_vars():
     config = build_mlflow_config()
 
     # Assert
-    assert config['host'] == 'http://test-host'
-    assert config['port'] == 8080
+    assert config['url'] == 'http://test-host:8080'
     assert config['username'] == 'test-user'
     assert config['password'] == 'test-pass'
 
@@ -28,8 +26,7 @@ def test_build_mlflow_config_with_env_vars():
 def test_build_mlflow_config_with_defaults():
     # Arrange
     # Clear any existing env vars
-    environ.pop('MLFLOW_HOST', None)
-    environ.pop('MLFLOW_PORT', None)
+    environ.pop('MLFLOW_URL', None)
     environ.pop('MLFLOW_USERNAME', None)
     environ.pop('MLFLOW_PASSWORD', None)
 
@@ -37,8 +34,7 @@ def test_build_mlflow_config_with_defaults():
     config = build_mlflow_config()
 
     # Assert
-    assert config['host'] == 'http://localhost'
-    assert config['port'] == 5080
+    assert config['url'] == 'http://localhost:5080'
     assert config['username'] == 'aignosi'
     assert config['password'] == 'aignosi'
 

tests/utils/test_exceptions.py

@@ -0,0 +1,79 @@
+"""Unit tests for custom exceptions with 100% coverage."""
+
+import pytest
+
+
+def test_model_training_error_with_default_message():
+    """Test ModelTrainingError with default message."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    error = ModelTrainingError(model_trained=True, model_saved=False)
+
+    assert error.model_trained is True
+    assert error.model_saved is False
+    assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=False)'
+
+
+def test_model_training_error_with_custom_message():
+    """Test ModelTrainingError with custom message."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    custom_msg = 'Custom error occurred during training'
+    error = ModelTrainingError(model_trained=False, model_saved=False, message=custom_msg)
+
+    assert error.model_trained is False
+    assert error.model_saved is False
+    assert str(error) == custom_msg
+
+
+def test_model_training_error_both_true():
+    """Test ModelTrainingError when both flags are True."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    error = ModelTrainingError(model_trained=True, model_saved=True)
+
+    assert error.model_trained is True
+    assert error.model_saved is True
+    assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=True)'
+
+
+def test_model_training_error_both_false():
+    """Test ModelTrainingError when both flags are False."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    error = ModelTrainingError(model_trained=False, model_saved=False)
+
+    assert error.model_trained is False
+    assert error.model_saved is False
+    assert str(error) == 'Model training workflow failed (model_trained=False, model_saved=False)'
+
+
+def test_model_training_error_is_exception():
+    """Test ModelTrainingError is an Exception subclass."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    error = ModelTrainingError(model_trained=True, model_saved=False)
+
+    assert isinstance(error, Exception)
+
+
+def test_model_training_error_can_be_raised():
+    """Test ModelTrainingError can be raised and caught."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    with pytest.raises(ModelTrainingError) as exc_info:
+        raise ModelTrainingError(model_trained=True, model_saved=False)
+
+    assert exc_info.value.model_trained is True
+    assert exc_info.value.model_saved is False
+
+
+def test_model_training_error_with_empty_message():
+    """Test ModelTrainingError with empty string message."""
+    from model_manager.utils.exceptions import ModelTrainingError
+
+    error = ModelTrainingError(model_trained=True, model_saved=True, message='')
+
+    assert error.model_trained is True
+    assert error.model_saved is True
+    assert str(error) == ''

tests/utils/test_logger_helper.py

@@ -0,0 +1,90 @@
+"""Unit tests for logger_helper module with 100% coverage."""
+
+from unittest.mock import MagicMock, patch
+
+
+@patch('model_manager.utils.logger_helper.SientiaLogger')
+def test_get_logger_creates_logger_instance(mock_sientia_logger):
+    """Test get_logger creates a SientiaLogger instance with the given name."""
+    from model_manager.utils.logger_helper import get_logger
+
+    mock_logger_instance = MagicMock()
+    mock_logger_instance.base_logger = MagicMock()
+    mock_sientia_logger.return_value = mock_logger_instance
+
+    result = get_logger('test_module')
+
+    mock_sientia_logger.assert_called_once_with('test_module')
+    assert result is mock_logger_instance
+
+
+@patch('model_manager.utils.logger_helper.SientiaLogger')
+def test_get_logger_disables_propagation(mock_sientia_logger):
+    """Test get_logger disables log propagation."""
+    from model_manager.utils.logger_helper import get_logger
+
+    mock_logger_instance = MagicMock()
+    mock_base_logger = MagicMock()
+    mock_base_logger.propagate = True
+    mock_logger_instance.base_logger = mock_base_logger
+    mock_sientia_logger.return_value = mock_logger_instance
+
+    get_logger('test_module')
+
+    assert mock_base_logger.propagate is False
+
+
+@patch('model_manager.utils.logger_helper.SientiaLogger')
+def test_get_logger_with_different_names(mock_sientia_logger):
+    """Test get_logger works with different logger names."""
+    from model_manager.utils.logger_helper import get_logger
+
+    mock_logger_instance = MagicMock()
+    mock_logger_instance.base_logger = MagicMock()
+    mock_sientia_logger.return_value = mock_logger_instance
+
+    logger1 = get_logger('module1')
+    logger2 = get_logger('module2')
+    logger3 = get_logger('my.nested.module')
+
+    assert mock_sientia_logger.call_count == 3
+    mock_sientia_logger.assert_any_call('module1')
+    mock_sientia_logger.assert_any_call('module2')
+    mock_sientia_logger.assert_any_call('my.nested.module')
+    assert logger1 is mock_logger_instance
+    assert logger2 is mock_logger_instance
+    assert logger3 is mock_logger_instance
+
+
+@patch('model_manager.utils.logger_helper.SientiaLogger')
+def test_get_logger_with_empty_name(mock_sientia_logger):
+    """Test get_logger with empty string name."""
+    from model_manager.utils.logger_helper import get_logger
+
+    mock_logger_instance = MagicMock()
+    mock_logger_instance.base_logger = MagicMock()
+    mock_sientia_logger.return_value = mock_logger_instance
+
+    result = get_logger('')
+
+    mock_sientia_logger.assert_called_once_with('')
+    assert result is mock_logger_instance
+    assert result.base_logger.propagate is False
+
+
+@patch('model_manager.utils.logger_helper.SientiaLogger')
+def test_get_logger_returns_configured_logger(mock_sientia_logger):
+    """Test get_logger returns the configured logger instance."""
+    from model_manager.utils.logger_helper import get_logger
+
+    mock_logger_instance = MagicMock()
+    mock_logger_instance.base_logger = MagicMock()
+    mock_logger_instance.base_logger.propagate = True
+    mock_sientia_logger.return_value = mock_logger_instance
+
+    result = get_logger('test_logger')
+
+    # Verify the logger is returned after configuration
+    assert result is mock_logger_instance
+    # Verify propagation was disabled
+    assert mock_logger_instance.base_logger.propagate is False

tests/worker/test_worker.py

@@ -1,494 +1,545 @@
-from unittest.mock import ANY, AsyncMock, MagicMock, patch
+"""Unit tests for worker module."""
+
+import asyncio
+import os
+import sys
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 
 
 @pytest.fixture
-def mock_env_vars(monkeypatch):
-    """Fixture to set up environment variables for tests."""
-    monkeypatch.setenv('POD_ID', 'test-pod-123')
-    monkeypatch.setenv('TEMPORAL_HOST', 'test-temporal:7233')
-    monkeypatch.setenv('TEMPORAL_NAMESPACE', 'test-namespace')
-    monkeypatch.setenv('HTTP_METRICS_PORT', '9090')
-    monkeypatch.setenv('HTTP_SDK_METRICS_PORT', '9091')
-    monkeypatch.setenv('PROJECT_NAME', 'test-project')
-    monkeypatch.setenv('POSTGRES_HOST', 'localhost')
-    monkeypatch.setenv('POSTGRES_PORT', '5432')
-    monkeypatch.setenv('POSTGRES_USER', 'test')
-    monkeypatch.setenv('POSTGRES_PASSWORD', 'test')
-    monkeypatch.setenv('POSTGRES_DBNAME', 'test')
-    monkeypatch.setenv('MLFLOW_HOST', 'http://localhost')
-    monkeypatch.setenv('MLFLOW_PORT', '5000')
-    monkeypatch.setenv('MLFLOW_USERNAME', 'test')
-    monkeypatch.setenv('MLFLOW_PASSWORD', 'test')
-    monkeypatch.setenv('MINIO_ENDPOINT_URL', 'http://localhost:9000')
-    monkeypatch.setenv('MINIO_ACCESS_KEY', 'test')
-    monkeypatch.setenv('MINIO_SECRET_KEY', 'test')
-    monkeypatch.setenv('MONGODB_USERNAME', 'test')
-    monkeypatch.setenv('MONGODB_PASSWORD', 'test')
-    monkeypatch.setenv('MONGODB_URL', 'localhost:27017')
-    monkeypatch.setenv('MONGODB_DATABASE_NAME', 'test')
+def mock_env_vars():
+    """Set up test environment variables."""
+    env_vars = {
+        'POD_ID': 'test-pod-123',
+        'HTTP_METRICS_PORT': '9090',
+        'HTTP_SDK_METRICS_PORT': '9091',
+        'TEMPORAL_HOST': 'localhost:7233',
+        'TEMPORAL_NAMESPACE': 'test-namespace',
+        'PROJECT_NAME': 'test-project',
+    }
+
+    with patch.dict(os.environ, env_vars, clear=False):
+        yield env_vars
 
 
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger."""
+    logger = Mock()
+    logger.custom_info = Mock()
+    logger.custom_error = Mock()
+    return logger
+
+
+@pytest.fixture
+def mock_temporal_client():
+    """Create a mock Temporal client."""
+    client_mock = AsyncMock()
+    client_mock.connect = AsyncMock()
+    return client_mock
+
+
+@pytest.fixture
+def mock_worker():
+    """Create a mock Temporal worker."""
+    worker_mock = Mock()
+    worker_mock.run = AsyncMock(return_value=None)
+    return worker_mock
+
+
+@pytest.fixture
+def mock_notification_handler():
+    """Create a mock notification handler."""
+    handler = Mock()
+    handler.shutdown = Mock()
+    return handler
+
+
+@pytest.fixture
+def mock_activities():
+    """Create a mock Activities instance."""
+    activities = AsyncMock()
+    activities.update_experiment_run = Mock()
+    activities.validate_train_params = Mock()
+    activities.train_model = Mock()
+    activities.cleanup_resources = Mock()
+    activities.shutdown = AsyncMock()
+    return activities
+
+
+def test_pod_id_from_env():
+    """Test that POD_ID is correctly read from environment."""
+    with patch.dict(os.environ, {'POD_ID': 'pod-test-123'}):
+        # Re-import to get new env value
+        import importlib
+
+        import model_manager.worker.worker as worker_module
+
+        importlib.reload(worker_module)
+
+        assert worker_module.POD_ID == 'pod-test-123'
+
+
+def test_sdk_metrics_port_default():
+    """Test that SDK_METRICS_PORT uses default value."""
+    with patch.dict(os.environ, {}, clear=True):
+        import importlib
+
+        import model_manager.worker.worker as worker_module
+
+        importlib.reload(worker_module)
+
+        assert worker_module.SDK_METRICS_PORT == 9091
+
+
+def test_sdk_metrics_port_from_env():
+    """Test that SDK_METRICS_PORT is read from environment."""
+    with patch.dict(os.environ, {'HTTP_SDK_METRICS_PORT': '8888'}):
+        import importlib
+
+        import model_manager.worker.worker as worker_module
+
+        importlib.reload(worker_module)
+
+        assert worker_module.SDK_METRICS_PORT == 8888
+
+
+@patch('model_manager.worker.worker.POD_ID', 'test-pod-123')
 @patch('model_manager.worker.worker.start_http_server')
 @patch('model_manager.worker.worker.metrics')
 def test_start_prometheus_server_success(mock_metrics, mock_start_http_server, mock_env_vars):
     """Test successful Prometheus server startup."""
-    # Import after patching to ensure mocks are in place
     from model_manager.worker.worker import start_prometheus_server
 
-    # Act
+    mock_app_up = Mock()
+    mock_metrics.APP_UP.labels.return_value = mock_app_up
+
     start_prometheus_server()
 
-    # Assert
+    # Verify HTTP server started
     mock_start_http_server.assert_called_once_with(9090)
+
+    # Verify APP_UP metric was set to 1
     mock_metrics.APP_UP.labels.assert_called_once_with(pod_id='test-pod-123')
-    mock_metrics.APP_UP.labels.return_value.set.assert_called_once_with(1)
+    mock_app_up.set.assert_called_once_with(1)
 
 
 @patch('model_manager.worker.worker.start_http_server')
-@patch('model_manager.worker.worker.os._exit')
-def test_start_prometheus_server_failure(mock_exit, mock_start_http_server, mock_env_vars):
-    """Test Prometheus server startup failure."""
-    # Arrange
-    mock_start_http_server.side_effect = Exception('Port already in use')
-
-    # Import after patching
+@patch('model_manager.worker.worker.metrics')
+def test_start_prometheus_server_custom_port(mock_metrics, mock_start_http_server):
+    """Test Prometheus server startup with custom port."""
     from model_manager.worker.worker import start_prometheus_server
 
-    # Act
+    with patch.dict(os.environ, {'HTTP_METRICS_PORT': '8080', 'POD_ID': 'custom-pod'}):
+        mock_app_up = Mock()
+        mock_metrics.APP_UP.labels.return_value = mock_app_up
+
+        start_prometheus_server()
+
+        mock_start_http_server.assert_called_once_with(8080)
+
+
+@patch('model_manager.worker.worker.start_http_server')
+@patch('model_manager.worker.worker.metrics')
+@patch('model_manager.worker.worker.os._exit')
+def test_start_prometheus_server_failure(
+    mock_exit, mock_metrics, mock_start_http_server, mock_env_vars
+):
+    """Test Prometheus server startup failure."""
+    from model_manager.worker.worker import start_prometheus_server
+
+    mock_start_http_server.side_effect = OSError('Port already in use')
+
     start_prometheus_server()
 
-    # Assert
-    mock_start_http_server.assert_called_once_with(9090)
+    # Verify exit was called with code 1
     mock_exit.assert_called_once_with(1)
 
 
 @pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
-@patch('model_manager.worker.worker.metrics')
-@patch('model_manager.worker.worker.asyncio.gather')
 @patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
+@patch('model_manager.worker.worker.client.Client')
 @patch('model_manager.worker.worker.Runtime')
 @patch('model_manager.worker.worker.Activities')
 @patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.build_mongodb_config')
+@patch('model_manager.worker.worker.build_postgres_config')
+@patch('model_manager.worker.worker.build_mlflow_config')
+@patch('model_manager.worker.worker.build_minio_config')
 @patch('model_manager.worker.worker.get_logger')
-async def test_main_success(
-    mock_get_logger,
+@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.metrics')
+async def test_main_successful_startup(
+    mock_metrics,
     mock_start_prometheus,
+    mock_get_logger,
+    mock_build_minio,
+    mock_build_mlflow,
+    mock_build_postgres,
+    mock_build_mongodb,
+    mock_notification_handler_class,
+    mock_activities_class,
+    mock_runtime_class,
+    mock_client_class,
+    mock_worker_class,
+    mock_env_vars,
+    mock_logger,
+    mock_temporal_client,
+    mock_worker,
     mock_notification_handler,
     mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_gather,
-    mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
 ):
-    """Test successful main function execution."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
-
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
-
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
-
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Mock gather to complete successfully
-    mock_gather.return_value = None
-
-    # Import and run
+    """Test successful main() execution until workers start."""
     from model_manager.worker.worker import main
 
-    # Act
-    await main()
+    # Setup mocks
+    mock_get_logger.return_value = mock_logger
+    mock_build_mongodb.return_value = {
+        'connection_string': 'mongodb://test',
+        'database_name': 'test_db',
+    }
+    mock_build_postgres.return_value = {}
+    mock_build_mlflow.return_value = {}
+    mock_build_minio.return_value = {}
 
-    # Assert
+    mock_notification_handler_class.return_value = mock_notification_handler
+    mock_activities_class.return_value = mock_activities
+
+    mock_runtime = Mock()
+    mock_runtime_class.return_value = mock_runtime
+
+    mock_client_instance = AsyncMock()
+    mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
+
+    mock_worker_instance = Mock()
+    mock_worker_instance.run = AsyncMock(
+        side_effect=asyncio.CancelledError()
+    )  # Simulate interruption
+    mock_worker_class.return_value = mock_worker_instance
+
+    mock_app_up = Mock()
+    mock_metrics.APP_UP.labels.return_value = mock_app_up
+
+    # Run main() and expect it to exit due to CancelledError
+    with pytest.raises(SystemExit) as exc_info:
+        await main()
+
+    assert exc_info.value.code == 1
+
+    # Verify all initialization steps were called
+    mock_get_logger.assert_called_once()
     mock_start_prometheus.assert_called_once()
-    mock_notification_handler.assert_called_once()
-    mock_activities.assert_called_once()
-    mock_client_connect.assert_called_once_with(
-        target_host='test-temporal:7233',
-        namespace='test-namespace',
-        runtime=ANY,
-    )
-    assert mock_worker.call_count == 1  # Only one worker created
-    mock_gather.assert_called_once()
-    mock_sys_exit.assert_called_once_with(1)
+    mock_notification_handler_class.assert_called_once()
+    mock_activities_class.assert_called_once()
+    mock_client_class.connect.assert_called_once()
+    mock_worker_class.assert_called_once()
+
+    # Verify cleanup was performed
+    mock_notification_handler.shutdown.assert_called_once()
+    mock_activities.shutdown.assert_called_once()
+    mock_app_up.set.assert_called_with(0)
 
 
 @pytest.mark.asyncio
-@patch('model_manager.worker.worker.asyncio.gather')
 @patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
+@patch('model_manager.worker.worker.client.Client')
 @patch('model_manager.worker.worker.Runtime')
 @patch('model_manager.worker.worker.Activities')
 @patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.build_mongodb_config')
+@patch('model_manager.worker.worker.build_postgres_config')
+@patch('model_manager.worker.worker.build_mlflow_config')
+@patch('model_manager.worker.worker.build_minio_config')
 @patch('model_manager.worker.worker.get_logger')
-@patch('model_manager.worker.worker.sys.exit')
+@patch('model_manager.worker.worker.start_prometheus_server')
 @patch('model_manager.worker.worker.metrics')
-async def test_main_exception_handling(
+async def test_main_handles_exception(
     mock_metrics,
-    mock_sys_exit,
-    mock_get_logger,
     mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_gather,
+    mock_get_logger,
+    mock_build_minio,
+    mock_build_mlflow,
+    mock_build_postgres,
+    mock_build_mongodb,
+    mock_notification_handler_class,
+    mock_activities_class,
+    mock_runtime_class,
+    mock_client_class,
+    mock_worker_class,
     mock_env_vars,
+    mock_logger,
 ):
-    """Test main function exception handling and cleanup."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
-
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
-
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
-
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Mock gather to raise an exception
-    mock_gather.side_effect = Exception('Worker failed')
-
-    # Import and run
+    """Test main() handles exceptions and performs cleanup."""
     from model_manager.worker.worker import main
 
-    # Act
-    await main()
+    # Setup mocks
+    mock_get_logger.return_value = mock_logger
+    mock_build_mongodb.return_value = {
+        'connection_string': 'mongodb://test',
+        'database_name': 'test_db',
+    }
+    mock_build_postgres.return_value = {}
+    mock_build_mlflow.return_value = {}
+    mock_build_minio.return_value = {}
 
-    # Assert - Verify cleanup was performed
+    mock_notification_handler = Mock()
+    mock_notification_handler.shutdown = Mock()
+    mock_notification_handler_class.return_value = mock_notification_handler
+
+    mock_activities = AsyncMock()
+    mock_activities.shutdown = AsyncMock()
+    mock_activities_class.return_value = mock_activities
+
+    mock_runtime = Mock()
+    mock_runtime_class.return_value = mock_runtime
+
+    mock_client_instance = AsyncMock()
+    mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
+
+    mock_worker_instance = Mock()
+    mock_worker_instance.run = AsyncMock(side_effect=RuntimeError('Worker failed'))
+    mock_worker_class.return_value = mock_worker_instance
+
+    mock_app_up = Mock()
+    mock_metrics.APP_UP.labels.return_value = mock_app_up
+
+    # Run main() and expect SystemExit
+    with pytest.raises(SystemExit) as exc_info:
+        await main()
+
+    assert exc_info.value.code == 1
+
+    # Verify error was logged
     mock_logger.custom_error.assert_called_once()
-    mock_handler.shutdown.assert_called_once()
-    mock_activities_instance.shutdown.assert_called_once()
-    mock_metrics.APP_UP.labels.assert_called_once_with(pod_id='test-pod-123')
-    mock_metrics.APP_UP.labels.return_value.set.assert_called_once_with(0)
-    mock_sys_exit.assert_called_once_with(1)
+    assert 'Worker failed' in str(mock_logger.custom_error.call_args)
+
+    # Verify cleanup was performed
+    mock_notification_handler.shutdown.assert_called_once()
+    mock_activities.shutdown.assert_called_once()
+    mock_app_up.set.assert_called_with(0)
 
 
 @pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
-@patch('model_manager.worker.worker.metrics')
 @patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
+@patch('model_manager.worker.worker.client.Client')
 @patch('model_manager.worker.worker.Runtime')
 @patch('model_manager.worker.worker.Activities')
 @patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.build_mongodb_config')
+@patch('model_manager.worker.worker.build_postgres_config')
+@patch('model_manager.worker.worker.build_mlflow_config')
+@patch('model_manager.worker.worker.build_minio_config')
 @patch('model_manager.worker.worker.get_logger')
-async def test_main_creates_only_one_worker(
-    mock_get_logger,
-    mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
+@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.metrics')
+async def test_main_temporal_client_configuration(
     mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
+    mock_start_prometheus,
+    mock_get_logger,
+    mock_build_minio,
+    mock_build_mlflow,
+    mock_build_postgres,
+    mock_build_mongodb,
+    mock_notification_handler_class,
+    mock_activities_class,
+    mock_runtime_class,
+    mock_client_class,
+    mock_worker_class,
+    mock_logger,
 ):
-    """Test that main creates only one worker with correct configurations."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
-
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
-
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
-
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Import
+    """Test that Temporal client is configured correctly."""
     from model_manager.worker.worker import main
 
-    # Mock gather to prevent infinite wait
-    with patch('model_manager.worker.worker.asyncio.gather', new_callable=AsyncMock):
-        # Act
+    with patch.dict(
+        os.environ,
+        {'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production'},
+    ):
+        # Setup mocks
+        mock_get_logger.return_value = mock_logger
+        mock_build_mongodb.return_value = {
+            'connection_string': 'mongodb://test',
+            'database_name': 'test_db',
+        }
+        mock_build_postgres.return_value = {}
+        mock_build_mlflow.return_value = {}
+        mock_build_minio.return_value = {}
+
+        mock_notification_handler = Mock()
+        mock_notification_handler.shutdown = Mock()
+        mock_notification_handler_class.return_value = mock_notification_handler
+
+        mock_activities = AsyncMock()
+        mock_activities.shutdown = AsyncMock()
+        mock_activities_class.return_value = mock_activities
+
+        mock_runtime = Mock()
+        mock_runtime_class.return_value = mock_runtime
+
+        mock_client_instance = AsyncMock()
+        mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
+
+        mock_worker_instance = Mock()
+        mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
+        mock_worker_class.return_value = mock_worker_instance
+
+        mock_app_up = Mock()
+        mock_metrics.APP_UP.labels.return_value = mock_app_up
+
+        # Run main()
+        with pytest.raises(SystemExit):
+            await main()
+
+        # Verify Temporal client was configured with correct parameters
+        mock_client_class.connect.assert_called_once_with(
+            target_host='temporal.example.com:7233', namespace='production', runtime=mock_runtime
+        )
+
+
+@pytest.mark.asyncio
+@patch('model_manager.worker.worker.Worker')
+@patch('model_manager.worker.worker.client.Client')
+@patch('model_manager.worker.worker.Runtime')
+@patch('model_manager.worker.worker.Activities')
+@patch('model_manager.worker.worker.NotificationHandler')
+@patch('model_manager.worker.worker.build_mongodb_config')
+@patch('model_manager.worker.worker.build_postgres_config')
+@patch('model_manager.worker.worker.build_mlflow_config')
+@patch('model_manager.worker.worker.build_minio_config')
+@patch('model_manager.worker.worker.get_logger')
+@patch('model_manager.worker.worker.start_prometheus_server')
+@patch('model_manager.worker.worker.metrics')
+async def test_main_worker_configuration(
+    mock_metrics,
+    mock_start_prometheus,
+    mock_get_logger,
+    mock_build_minio,
+    mock_build_mlflow,
+    mock_build_postgres,
+    mock_build_mongodb,
+    mock_notification_handler_class,
+    mock_activities_class,
+    mock_runtime_class,
+    mock_client_class,
+    mock_worker_class,
+    mock_env_vars,
+    mock_logger,
+):
+    """Test that Temporal worker is configured with correct parameters."""
+    from model_manager.worker.worker import main
+
+    # Setup mocks
+    mock_get_logger.return_value = mock_logger
+    mock_build_mongodb.return_value = {
+        'connection_string': 'mongodb://test',
+        'database_name': 'test_db',
+    }
+    mock_build_postgres.return_value = {}
+    mock_build_mlflow.return_value = {}
+    mock_build_minio.return_value = {}
+
+    mock_notification_handler = Mock()
+    mock_notification_handler_class.return_value = mock_notification_handler
+
+    mock_activities = AsyncMock()
+    mock_activities.update_experiment_run = Mock()
+    mock_activities.validate_train_params = Mock()
+    mock_activities.train_model = Mock()
+    mock_activities.cleanup_resources = Mock()
+    mock_activities.shutdown = AsyncMock()
+    mock_activities_class.return_value = mock_activities
+
+    mock_runtime = Mock()
+    mock_runtime_class.return_value = mock_runtime
+
+    mock_client_instance = AsyncMock()
+    mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
+
+    mock_worker_instance = Mock()
+    mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
+    mock_worker_class.return_value = mock_worker_instance
+
+    mock_app_up = Mock()
+    mock_metrics.APP_UP.labels.return_value = mock_app_up
+
+    # Run main()
+    with pytest.raises(SystemExit):
         await main()
 
-    # Assert - Verify only one worker was created
-    assert mock_worker.call_count == 1
+    # Verify Worker was created with correct configuration
+    mock_worker_class.assert_called_once()
+    call_args = mock_worker_class.call_args
 
-    # Verify worker (train_model-queue)
-    first_call = mock_worker.call_args_list[0]
-    assert first_call[1]['task_queue'] == 'train_model-queue'
-    assert 'TrainModel' in str(first_call[1]['workflows'])
+    assert call_args[0][0] == mock_client_instance  # temporal_client
+    assert call_args[1]['task_queue'] == 'train_model-queue'
+    assert call_args[1]['max_concurrent_workflow_tasks'] == 50
+    assert call_args[1]['max_concurrent_activities'] == 50
+    assert call_args[1]['max_concurrent_local_activities'] == 50
+    assert call_args[1]['max_cached_workflows'] == 200
+
+    # Verify activities are included
+    activities_list = call_args[1]['activities']
+    assert mock_activities.update_experiment_run in activities_list
+    assert mock_activities.validate_train_params in activities_list
+    assert mock_activities.train_model in activities_list
+    assert mock_activities.cleanup_resources in activities_list
 
 
-@pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
+@patch('model_manager.worker.worker.asyncio.run')
+def test_main_entrypoint(mock_asyncio_run):
+    """Test the __main__ entrypoint."""
+    # Import and execute the main block
+    with patch.object(sys, 'argv', ['worker.py']):
+        import model_manager.worker.worker as worker_module
+
+        # Simulate running the module
+        worker_module.main = AsyncMock()
+
+        # This would normally be called by asyncio.run(main())
+        # We just verify the pattern is correct
+        assert callable(worker_module.main)
+
+
+def test_worker_module_docstring():
+    """Test that worker module has comprehensive documentation."""
+    import model_manager.worker.worker as worker_module
+
+    assert worker_module.__doc__ is not None
+    assert 'Temporal' in worker_module.__doc__
+    assert 'worker' in worker_module.__doc__
+
+
+@patch('model_manager.worker.worker.start_http_server')
 @patch('model_manager.worker.worker.metrics')
-@patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
-@patch('model_manager.worker.worker.Runtime')
-@patch('model_manager.worker.worker.Activities')
-@patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
-@patch('model_manager.worker.worker.get_logger')
-async def test_main_initializes_activities_with_configs(
-    mock_get_logger,
-    mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
+def test_start_prometheus_server_prints_success(
+    mock_metrics, mock_start_http_server, capsys, mock_env_vars
 ):
-    """Test that main initializes Activities with correct configurations."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
+    """Test that start_prometheus_server prints success message."""
+    from model_manager.worker.worker import start_prometheus_server
 
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
+    mock_app_up = Mock()
+    mock_metrics.APP_UP.labels.return_value = mock_app_up
 
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
+    start_prometheus_server()
 
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Import
-    from model_manager.worker.worker import main
-
-    # Mock gather to prevent infinite wait
-    with patch('model_manager.worker.worker.asyncio.gather', new_callable=AsyncMock):
-        # Act
-        await main()
-
-    # Assert - Verify Activities was initialized with correct parameters
-    mock_activities.assert_called_once()
-    call_kwargs = mock_activities.call_args[1]
-    assert 'postgres_config' in call_kwargs
-    assert 'mlflow_config' in call_kwargs
-    assert 'minio_config' in call_kwargs
-    assert call_kwargs['logger'] == mock_logger
-    assert call_kwargs['notification_handler'] == mock_handler
+    captured = capsys.readouterr()
+    assert 'Prometheus server started on port 9090' in captured.out
 
 
-@pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
+@patch('model_manager.worker.worker.start_http_server')
 @patch('model_manager.worker.worker.metrics')
-@patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
-@patch('model_manager.worker.worker.Runtime')
-@patch('model_manager.worker.worker.Activities')
-@patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
-@patch('model_manager.worker.worker.get_logger')
-async def test_main_uses_environment_variables(
-    mock_get_logger,
-    mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
+@patch('model_manager.worker.worker.os._exit')
+def test_start_prometheus_server_prints_failure(
+    mock_exit, mock_metrics, mock_start_http_server, capsys, mock_env_vars
 ):
-    """Test that main uses environment variables correctly."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
+    """Test that start_prometheus_server prints failure message."""
+    from model_manager.worker.worker import start_prometheus_server
 
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
+    mock_start_http_server.side_effect = Exception('Test error')
 
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
+    start_prometheus_server()
 
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Import
-    from model_manager.worker.worker import main
-
-    # Mock gather to prevent infinite wait
-    with patch('model_manager.worker.worker.asyncio.gather', new_callable=AsyncMock):
-        # Act
-        await main()
-
-    # Assert - Verify environment variables were used
-    mock_client_connect.assert_called_once_with(
-        target_host='test-temporal:7233',
-        namespace='test-namespace',
-        runtime=ANY,
-    )
-
-    mock_notification_handler.assert_called_once()
-    notification_call_kwargs = mock_notification_handler.call_args[1]
-    assert notification_call_kwargs['project_name'] == 'test-project'
-
-
-@pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
-@patch('model_manager.worker.worker.metrics')
-@patch('model_manager.worker.worker.asyncio.gather')
-@patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
-@patch('model_manager.worker.worker.Runtime')
-@patch('model_manager.worker.worker.Activities')
-@patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
-@patch('model_manager.worker.worker.get_logger')
-async def test_main_cleanup_with_none_notification_handler(
-    mock_get_logger,
-    mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_gather,
-    mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
-):
-    """Test cleanup when notification_handler is None (line 194 branch False)."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
-
-    # Return None for notification_handler
-    mock_notification_handler.return_value = None
-
-    mock_activities_instance = MagicMock()
-    mock_activities_instance.shutdown = AsyncMock()
-    mock_activities.return_value = mock_activities_instance
-
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Mock gather to complete
-    mock_gather.return_value = None
-
-    # Import and run
-    from model_manager.worker.worker import main
-
-    # Act
-    await main()
-
-    # Assert - notification_handler.shutdown() should NOT be called (line 194 False)
-    # Since notification_handler is None, we can't call shutdown on it
-    mock_activities_instance.shutdown.assert_called_once()
-    mock_sys_exit.assert_called_once_with(1)
-
-
-@pytest.mark.asyncio
-@patch('model_manager.worker.worker.sys.exit')
-@patch('model_manager.worker.worker.metrics')
-@patch('model_manager.worker.worker.asyncio.gather')
-@patch('model_manager.worker.worker.Worker')
-@patch('model_manager.worker.worker.client.Client.connect')
-@patch('model_manager.worker.worker.Runtime')
-@patch('model_manager.worker.worker.Activities')
-@patch('model_manager.worker.worker.NotificationHandler')
-@patch('model_manager.worker.worker.start_prometheus_server')
-@patch('model_manager.worker.worker.get_logger')
-async def test_main_cleanup_with_falsy_activities(
-    mock_get_logger,
-    mock_start_prometheus,
-    mock_notification_handler,
-    mock_activities,
-    mock_runtime,
-    mock_client_connect,
-    mock_worker,
-    mock_gather,
-    mock_metrics,
-    mock_sys_exit,
-    mock_env_vars,
-):
-    """Test cleanup when activities evaluates to False (line 196 branch False)."""
-    # Arrange
-    mock_logger = MagicMock()
-    mock_get_logger.return_value = mock_logger
-
-    mock_handler = MagicMock()
-    mock_notification_handler.return_value = mock_handler
-
-    # Create a falsy activities object (empty list, 0, False, etc.)
-    # Using an object that evaluates to False but doesn't cause AttributeError
-    class FalsyActivities:
-        def __bool__(self):
-            return False
-
-        def __getattr__(self, name):
-            # Return mock methods to avoid AttributeError during worker creation
-            return MagicMock()
-
-    falsy_activities = FalsyActivities()
-    mock_activities.return_value = falsy_activities
-
-    mock_temporal_client = AsyncMock()
-    mock_client_connect.return_value = mock_temporal_client
-
-    mock_worker_instance = MagicMock()
-    mock_worker_instance.run = MagicMock(return_value=AsyncMock())
-    mock_worker.return_value = mock_worker_instance
-
-    # Mock gather to complete
-    mock_gather.return_value = None
-
-    # Import and run
-    from model_manager.worker.worker import main
-
-    # Act
-    await main()
-
-    # Assert - notification_handler.shutdown() is called, but activities.shutdown() is NOT
-    mock_handler.shutdown.assert_called_once()
-    # activities is falsy, so shutdown should NOT be called
-    mock_sys_exit.assert_called_once_with(1)
+    captured = capsys.readouterr()
+    assert 'Failed to start Prometheus server' in captured.out
+    assert 'Test error' in captured.out

todo-list.txt

@@ -0,0 +1,2 @@
+- remover os testes das classes alteradas e refazer de novo depois
+- no final alterar o readme

validate.sh

@@ -1,9 +1,55 @@
 #!/bin/bash
 # Model Manager Code Validation Script
 # This script runs all code quality checks before committing or deploying
+#
+# Usage:
+#   ./validate.sh              # Run all checks including tests (default)
+#   ./validate.sh --no-tests   # Skip unit tests
+#   ./validate.sh --skip-tests # Skip unit tests (alias)
+#   ./validate.sh --only-tests # Run only unit tests
+#   ./validate.sh --fix        # Auto-fix formatting and linting, then run validations (no tests)
 
 set -e  # Exit on any error
 
+# Parse command line arguments
+RUN_TESTS=true
+ONLY_TESTS=false
+FIX_MODE=false
+
+for arg in "$@"; do
+    case $arg in
+        --no-tests|--skip-tests)
+            RUN_TESTS=false
+            shift
+            ;;
+        --only-tests)
+            ONLY_TESTS=true
+            shift
+            ;;
+        --fix)
+            FIX_MODE=true
+            RUN_TESTS=false
+            shift
+            ;;
+        --help|-h)
+            echo "Usage: $0 [OPTIONS]"
+            echo ""
+            echo "Options:"
+            echo "  --no-tests, --skip-tests    Skip unit tests (default: run tests)"
+            echo "  --only-tests                Run only unit tests"
+            echo "  --fix                       Auto-fix formatting and linting, then run validations (no tests)"
+            echo "  --help, -h                  Show this help message"
+            echo ""
+            exit 0
+            ;;
+        *)
+            echo "Unknown option: $arg"
+            echo "Use --help for usage information"
+            exit 1
+            ;;
+    esac
+done
+
 # Colors for output
 RED='\033[0;31m'
 GREEN='\033[0;32m'
@@ -16,6 +62,40 @@ echo -e "${BLUE}║  Model Manager - Code Validation Suite                ║${N
 echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
 echo ""
 
+# Handle --fix mode
+if [ "$FIX_MODE" = true ]; then
+    echo -e "${BLUE}🔧 Running auto-fix mode...${NC}"
+    echo ""
+    
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    echo -e "${BLUE}▶ Auto-fixing code formatting (Ruff)${NC}"
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    ruff format model_manager/ tests/
+    echo -e "${GREEN}✅ Code formatting applied${NC}"
+    echo ""
+    
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    echo -e "${BLUE}▶ Auto-fixing linting issues (Ruff)${NC}"
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    ruff check --fix model_manager/ tests/
+    echo -e "${GREEN}✅ Linting fixes applied${NC}"
+    echo ""
+    
+    echo -e "${YELLOW}ℹ️  Now running validations (without tests)...${NC}"
+    echo ""
+fi
+
+# Handle --only-tests mode
+if [ "$ONLY_TESTS" = true ]; then
+    echo -e "${BLUE}🧪 Running only unit tests...${NC}"
+    echo ""
+fi
+
+if [ "$RUN_TESTS" = false ] && [ "$ONLY_TESTS" = false ]; then
+    echo -e "${YELLOW}ℹ️  Unit tests will be skipped${NC}"
+    echo ""
+fi
+
 # Check if virtual environment is activated
 if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then
     echo -e "${YELLOW}⚠️  Warning: No virtual environment detected${NC}"
@@ -46,29 +126,45 @@ run_step() {
 # Track failures
 FAILED_STEPS=()
 
-# Step 1: Code Formatting Check (Ruff)
-if ! run_step "1. Code Formatting (Ruff)" "ruff format --check model_manager/ tests/"; then
-    FAILED_STEPS+=("Code Formatting")
-fi
+# Handle --only-tests mode
+if [ "$ONLY_TESTS" = true ]; then
+    # Step 5: Unit Tests (pytest)
+    if ! run_step "Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
+        FAILED_STEPS+=("Unit Tests")
+    fi
+else
+    # Step 1: Code Formatting Check (Ruff)
+    if ! run_step "1. Code Formatting (Ruff)" "ruff format --check model_manager/ tests/"; then
+        FAILED_STEPS+=("Code Formatting")
+    fi
 
-# Step 2: Linting (Ruff)
-if ! run_step "2. Code Linting (Ruff)" "ruff check model_manager/ tests/"; then
-    FAILED_STEPS+=("Linting")
-fi
+    # Step 2: Linting (Ruff)
+    if ! run_step "2. Code Linting (Ruff)" "ruff check model_manager/ tests/"; then
+        FAILED_STEPS+=("Linting")
+    fi
 
-# Step 3: Type Checking (mypy)
-if ! run_step "3. Type Checking (mypy)" "mypy model_manager/"; then
-    FAILED_STEPS+=("Type Checking")
-fi
+    # Step 3: Type Checking (mypy)
+    if ! run_step "3. Type Checking (mypy)" "mypy model_manager/"; then
+        FAILED_STEPS+=("Type Checking")
+    fi
 
-# Step 4: Security Analysis (Bandit)
-if ! run_step "4. Security Analysis (Bandit)" "bandit -r model_manager/ -ll -q"; then
-    FAILED_STEPS+=("Security Analysis")
-fi
+    # Step 4: Security Analysis (Bandit)
+    if ! run_step "4. Security Analysis (Bandit)" "bandit -r model_manager/ -ll -q"; then
+        FAILED_STEPS+=("Security Analysis")
+    fi
 
-# Step 5: Unit Tests (pytest)
-if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
-    FAILED_STEPS+=("Unit Tests")
+    # Step 5: Unit Tests (pytest)
+    if [ "$RUN_TESTS" = true ]; then
+        if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
+            FAILED_STEPS+=("Unit Tests")
+        fi
+    else
+        echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+        echo -e "${BLUE}▶ 5. Unit Tests (pytest)${NC}"
+        echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+        echo -e "${YELLOW}⏭️  Unit Tests - SKIPPED${NC}"
+        echo ""
+    fi
 fi
 
 # Summary