SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.

.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"

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,6 +88,7 @@ class ExperimentTracking(Postgres):
             logger=logger,
             notification_handler=notification_handler,
         )
+
         self.logger = logger
         self.notification_handler = notification_handler
 
@@ -105,6 +109,14 @@ class ExperimentTracking(Postgres):
                 # Silently ignore errors during garbage collection
                 pass
 
+    async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
+        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)
+
     @activity.defn(name='update_experiment_run')
     async def update_experiment_run(self, input_data: dict[str, Any]) -> None:
         """
@@ -128,32 +140,6 @@ 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']
@@ -164,71 +150,86 @@ class ExperimentTracking(Postgres):
 
         try:
             self.info(
-                f'Updating experiment run {experiment_run_id} with type: {update_type}', metadata
+                f'Updating experiment run {experiment_run_id} with status: {status}', metadata
             )
 
-            # Validate parameters based on update type
-            query_params: tuple[Any, ...]
+            query_params: dict[str, Any]
+
             if update_type == UpdateType.STATUS:
-                if not status:
+                if not isinstance(status, str) or 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
+                    SET status = :status, updated_at = :updated_at
+                    WHERE id = :experiment_run_id
                 """
-                query_params = (status, datetime.now(UTC), experiment_run_id)
 
+                query_params = {
+                    'status': status,
+                    'updated_at': datetime.now(UTC),
+                    'experiment_run_id': 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'
-                    )
+                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 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
+                    SET status = :status, error_message = :error_message, updated_at = :updated_at
+                    WHERE id = :experiment_run_id
                 """
-                query_params = (
-                    status,
-                    error_message,
-                    datetime.now(UTC),
-                    experiment_run_id,
-                )
 
+                query_params = {
+                    'status': status,
+                    'error_message': error_message,
+                    'updated_at': datetime.now(UTC),
+                    'experiment_run_id': experiment_run_id,
+                }
             elif update_type == UpdateType.MODEL_SAVED:
-                if not run_name:
+                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 = %s, status = %s, updated_at = %s
-                    WHERE id = %s
+                    SET run_name = :run_name, status = :status, updated_at = :updated_at
+                    WHERE id = :experiment_run_id
                 """
-                query_params = (run_name, status, datetime.now(UTC), experiment_run_id)
 
+                query_params = {
+                    'run_name': run_name,
+                    'status': status,
+                    'updated_at': datetime.now(UTC),
+                    'experiment_run_id': 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 +238,6 @@ 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,12 @@ 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 +92,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',
@@ -117,14 +105,11 @@ class Training(BaseActivity):
                 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.
 
@@ -156,51 +141,42 @@ class Training(BaseActivity):
             # Returns: TrainModelResult(...)
         """
         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 +186,39 @@ class Training(BaseActivity):
                 attachment_content=trace,
             )
 
-            # Log error with metadata
             self.error(trace, metadata=metadata)
 
-            # Re-raise exception to stop workflow
+            raise ModelTrainingError(
+                model_trained=model_trained,
+                model_saved=model_saved,
+            ) from e
+
+    @activity.defn(name='cleanup_resources')
+    async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
+        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,
+            )
+
+            self.error(trace, metadata=metadata)
             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

model_manager/utils/connectors_config.py

@@ -51,8 +51,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,12 @@
+class ModelTrainingError(Exception):
+    def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None):
+        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/models/train_model_params.py

@@ -184,22 +184,22 @@ class TrainModelParams:
             >>> 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 +218,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/repository/model_repository.py

@@ -9,6 +9,7 @@ and logging model runs to MLFlow.
 
 """
 
+import os
 import shutil
 from datetime import datetime
 from os import makedirs, path
@@ -22,14 +23,81 @@ from model_manager.sientia.reports import Reports  # type: ignore[import-untyped
 from model_manager.utils.models.train_model_result import TrainModelResult
 
 
-class MLFlowRepository:
-    def __init__(self, host, username, password, logger: Logger):
-        self.model_serving = ModelServing(
-            tracking_uri=host, username=username, password=password, logger=logger
-        )
+class ModelRepository:
+    def __init__(self, url, username, password, logger: Logger):
+        self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
+
         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
+        self.logger.info(f'Starting model save for experiment: {experiment_name}')
+
+        # Step 1: Generate next run name
+        self.logger.info('Generating run name')
+        train_result.run_name = self._get_next_run_name(experiment_name)
+        self.logger.info(f'Generated run name: {train_result.run_name}')
+
+        # Step 2: Generate artifacts (reports, CSV files)
+        self.logger.info('Generating artifacts')
+        train_result = self._generate_artifacts(train_result)
+        self.logger.info('Artifacts generated successfully')
+
+        # Step 3: Save run to MLflow
+        self.logger.info('Saving run to MLflow')
+        self._save_run(train_result)
+
+        self.logger.info(
+            f'Model saved successfully - Run: {train_result.run_name}, '
+            f'Experiment: {experiment_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
+
+        self.logger.info(f'Cleaning up run directory: {run_dir}')
+
+        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 +114,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 +155,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 +190,48 @@ 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}"
-                )
-
-        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 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)
 
     def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
         """
@@ -258,7 +314,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 +353,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 +412,16 @@ 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}')
 
             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,129 @@
+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.).
+        """
+        self.logger.info(f'Fetching file from MinIO: {bucket_name}/{file_name}')
+        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.logger.info(f'Deleting file from MinIO: {bucket_name}/{file_name}')
+        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,17 @@ class TrainingRepository:
             ValueError: If transformed data is empty
             Exception: If data loading, preprocessing, or training fails
         """
-        # Load data from BytesIO
+        self.logger.info('Loading data from BytesIO file')
         data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
-
-        # Initialize and fit data preprocessor
-        process_data = self.init_data_preprocessor(params)
+        self.logger.info('Initializing and fitting data preprocessor')
+        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
+        self.logger.info('Splitting 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 +85,18 @@ class TrainingRepository:
             random_state=42,
         )
 
-        # Prepare training data
+        self.logger.info('Preparing 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
+        self.logger.info('Training linear regression model')
         regr = LinearRegressionModel(
             target_variable=params.target_variable,
             variable_columns=params.variable_columns,
         )
+
         regr.fit(data_train)
 
-        # Return training result
         return TrainModelResult(
             params=params,
             process_data=process_data,
@@ -110,7 +108,70 @@ 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)
+        """
+        self.logger.info('Making predictions on test set')
+        y_pred_array = tmr.regr.predict(tmr.x_test)
+
+        if params.use_scaler:
+            scaler = tmr.process_data.get_scaler()
+            self.logger.info('Denormalizing 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)
+
+            self.logger.info('Denormalizing 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)
+
+        self.logger.info('Adding 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'
+
+        self.logger.info('Reordering 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()
+
+        self.logger.info('Calculating 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_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
         """
         Initialize dictionary containing scaling parameters for features and target.
 
@@ -152,69 +213,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.
 

model_manager/worker/worker.py

@@ -128,18 +128,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,

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,14 +53,6 @@ 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,
-    )
-
 
 @workflow.defn(name='train_model')
 class TrainModel:
@@ -113,8 +94,6 @@ 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 = {
@@ -124,33 +103,21 @@ class TrainModel:
             }
         }
 
-        # 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 +141,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 +174,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 +193,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 +228,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 +293,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 +352,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)

scripts/run_training_test.py

@@ -0,0 +1,227 @@
+#!/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(
+        ['mc', 'cp', str(source_path), target_uri],
+        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

@@ -13,7 +13,7 @@ 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"
+    "55432: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"

tests/activities/test_activities.py

@@ -1,535 +0,0 @@
-from unittest.mock import ANY, MagicMock, patch
-
-from pytest import mark
-
-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
-
-
-@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 = {
-        '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,
-    )
-
-    assert isinstance(activities, Activities)
-    assert isinstance(activities, ExperimentTracking)
-    assert isinstance(activities, MLFlow)
-    assert isinstance(activities, Training)
-
-    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_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,
-        endpoint_url=minio_config['endpoint_url'],
-        access_key=minio_config['access_key'],
-        secret_key=minio_config['secret_key'],
-        region=minio_config['region'],
-        use_ssl=minio_config['use_ssl'],
-        max_retry_attempts=minio_config['max_retry_attempts'],
-        retry_mode=minio_config['retry_mode'],
-        connect_timeout=minio_config['connect_timeout'],
-        read_timeout=minio_config['read_timeout'],
-        logger=logger,
-        notification_handler=notification_handler,
-    )
-
-    mock_training_init.assert_called_once_with(
-        ANY, logger=logger, notification_handler=notification_handler
-    )
-
-
-@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(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """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()
-
-    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()
-
-    # 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()
-
-
-@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(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """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()
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        minio_config=minio_config,
-        logger=logger,
-        notification_handler=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__()
-
-        # 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(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """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()
-
-    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()
-
-    # Mock super().__del__ to raise an exception
-    mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed'))
-
-    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
-
-
-@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__()
-        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'
-
-
-@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(
-    mock_training_init,
-    mock_minio_init,
-    mock_mlflow_init,
-    mock_experiment_tracking_init,
-):
-    """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
-
-    activities = Activities(
-        postgres_config=postgres_config,
-        mlflow_config=mlflow_config,
-        minio_config=minio_config,
-        logger=logger,
-        notification_handler=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."""
-
-        pass
-
-    # Patch super() to return an instance that doesn't have __del__
-    mock_super_instance = MockSuperWithoutDel()
-
-    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

tests/activities/test_experiment_tracking.py

@@ -1,564 +0,0 @@
-"""Unit tests for ExperimentTracking activity."""
-
-from unittest.mock import AsyncMock, MagicMock, patch
-
-from pytest import mark, raises
-
-from model_manager.activities.experiment_tracking import ExperimentTracking, UpdateType
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_status_success(mock_postgres_init):
-    """Test successful status update."""
-    mock_postgres_init.return_value = None
-
-    # Create instance
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Mock methods
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})
-
-    # Test data
-    input_data = {
-        'metadata': {'workflow_id': 'test-123'},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS,
-        'status': 'TRAINING_SUCCESS',
-    }
-
-    # Execute
-    await tracking.update_experiment_run(input_data)
-
-    # Assertions
-    tracking.info.assert_called()
-    tracking.execute_query.assert_called_once()
-    call_args = tracking.execute_query.call_args
-    assert 'UPDATE experiment_run' in call_args[0][0]
-    assert 'SET status = %s' in call_args[0][0]
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_status_with_error_success(mock_postgres_init):
-    """Test successful status update with error message."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-123'},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS_WITH_ERROR,
-        'status': 'TRAINING_ERROR',
-        'error_message': 'Model training failed due to insufficient data',
-    }
-
-    await tracking.update_experiment_run(input_data)
-
-    tracking.info.assert_called()
-    tracking.execute_query.assert_called_once()
-    call_args = tracking.execute_query.call_args
-    assert 'UPDATE experiment_run' in call_args[0][0]
-    assert 'SET status = %s, error_message = %s' in call_args[0][0]
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_error_message_truncation(mock_postgres_init):
-    """Test that error messages longer than 1024 chars are truncated."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})
-
-    # Create error message longer than 1024 characters
-    long_error = 'A' * 2000
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS_WITH_ERROR,
-        'status': 'TRAINING_ERROR',
-        'error_message': long_error,
-    }
-
-    await tracking.update_experiment_run(input_data)
-
-    # Check that error message was truncated to 1024 chars
-    call_args = tracking.execute_query.call_args
-    query_params = call_args[0][1]
-    assert len(query_params[1]) == 1024  # error_message is second param
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_model_saved_success(mock_postgres_init):
-    """Test successful model saved update."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(return_value={'rowcount': 1})
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.MODEL_SAVED,
-        'run_name': 'experiment-model-123',
-        'status': 'MLFLOW_SENT',
-    }
-
-    await tracking.update_experiment_run(input_data)
-
-    tracking.info.assert_called()
-    tracking.execute_query.assert_called_once()
-    call_args = tracking.execute_query.call_args
-    assert 'UPDATE experiment_run' in call_args[0][0]
-    assert 'SET run_name = %s, status = %s' in call_args[0][0]
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_missing_status_raises_error(mock_postgres_init):
-    """Test that missing status parameter raises ValueError."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS,
-        # Missing 'status' parameter
-    }
-
-    with raises(RuntimeError, match='Error updating experiment run'):
-        await tracking.update_experiment_run(input_data)
-
-    tracking.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_missing_error_message_raises_error(mock_postgres_init):
-    """Test that missing error_message parameter raises ValueError."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS_WITH_ERROR,
-        'status': 'TRAINING_ERROR',
-        # Missing 'error_message' parameter
-    }
-
-    with raises(RuntimeError, match='Error updating experiment run'):
-        await tracking.update_experiment_run(input_data)
-
-    tracking.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_missing_run_name_raises_error(mock_postgres_init):
-    """Test that missing run_name parameter raises ValueError."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.MODEL_SAVED,
-        # Missing 'run_name' parameter
-    }
-
-    with raises(RuntimeError, match='Error updating experiment run'):
-        await tracking.update_experiment_run(input_data)
-
-    tracking.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_invalid_update_type_raises_error(mock_postgres_init):
-    """Test that invalid update_type raises ValueError."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 456,
-        'update_type': 'INVALID_TYPE',
-        'status': 'TRAINING_SUCCESS',
-    }
-
-    with raises(RuntimeError, match='Error updating experiment run'):
-        await tracking.update_experiment_run(input_data)
-
-    tracking.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_no_rows_updated_raises_error(mock_postgres_init):
-    """Test that zero rows updated raises ValueError."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(return_value={'rowcount': 0})
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {},
-        'experiment_run_id': 999,  # Non-existent ID
-        'update_type': UpdateType.STATUS,
-        'status': 'TRAINING_SUCCESS',
-    }
-
-    with raises(RuntimeError, match='Error updating experiment run'):
-        await tracking.update_experiment_run(input_data)
-
-    tracking.send_notification.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-async def test_update_experiment_run_sends_notification_on_error(mock_postgres_init):
-    """Test that notification is sent when update fails."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    tracking.info = MagicMock()
-    tracking.execute_query = AsyncMock(side_effect=Exception('Database error'))
-    tracking.send_notification = MagicMock()
-    tracking.error = MagicMock()
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-123'},
-        'experiment_run_id': 456,
-        'update_type': UpdateType.STATUS,
-        'status': 'TRAINING_SUCCESS',
-    }
-
-    with raises(RuntimeError):
-        await tracking.update_experiment_run(input_data)
-
-    # Verify notification was sent
-    tracking.send_notification.assert_called_once()
-    call_args = tracking.send_notification.call_args
-    assert call_args[1]['notification_id'] == 'UPDATE_EXPERIMENT_RUN_ERROR'
-    assert call_args[1]['metadata'] == {'workflow_id': 'test-123'}
-
-
-def test_update_type_enum_values():
-    """Test UpdateType enum has correct values."""
-    assert UpdateType.STATUS == 'status'
-    assert UpdateType.STATUS_WITH_ERROR == 'status_with_error'
-    assert UpdateType.MODEL_SAVED == 'model_saved'
-
-
-# Tests for __del__ method - 100% coverage
-
-
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-def test___del___with_engine_and_parent_del_exists(mock_postgres_init):
-    """Test __del__ calls parent destructor when engine exists and parent has __del__."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Add engine attribute to simulate Postgres initialization
-    tracking.engine = MagicMock()
-
-    # Track if parent __del__ was called
-    parent_del_called = []
-
-    def mock_parent_del(self):
-        """Mock parent __del__ that tracks when it's called."""
-        parent_del_called.append(True)
-
-    # Patch parent class to have __del__ method
-    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
-        # Trigger __del__ - should call parent __del__ (line 103)
-        tracking.__del__()
-
-        # Verify parent __del__ was called (line 103 executed)
-        assert len(parent_del_called) == 1, 'Parent __del__ should have been called once'
-
-
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-def test___del___without_engine(mock_postgres_init):
-    """Test __del__ does not call parent destructor when engine attribute is missing."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Ensure engine attribute does not exist
-    if hasattr(tracking, 'engine'):
-        delattr(tracking, 'engine')
-
-    # Mock parent __del__ to track if it's called
-    mock_parent_del = MagicMock()
-
-    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
-        # Trigger __del__ - should NOT call parent __del__ (line 100 is False)
-        tracking.__del__()
-
-        # Verify parent __del__ was NOT called
-        mock_parent_del.assert_not_called()
-
-
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-def test___del___when_parent_has_no_del_method(mock_postgres_init):
-    """Test __del__ handles case when parent class has no __del__ method - covers line 102 false branch."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Add engine attribute to pass the first hasattr check (line 100)
-    tracking.engine = MagicMock()
-
-    # Create a mock class without __del__ method
-    class MockSuperWithoutDel:
-        """Mock class that explicitly does not have __del__ method."""
-
-        pass
-
-    # Patch super() to return an instance without __del__
-    mock_super_instance = MockSuperWithoutDel()
-
-    with patch('builtins.super', return_value=mock_super_instance):
-        # Trigger __del__ - should handle the case when hasattr(super(), '__del__') is False (line 102)
-        try:
-            tracking.__del__()
-            # Test passes - the false branch of line 102 was executed without error
-        except Exception as e:
-            raise AssertionError(
-                f'__del__ should handle super() without __del__ method, but raised: {e}'
-            ) from e
-
-
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-def test___del___handles_exception_from_parent_del(mock_postgres_init):
-    """Test __del__ handles exceptions from parent destructor gracefully - covers line 104."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Add engine attribute
-    tracking.engine = MagicMock()
-
-    # Mock parent __del__ to raise an exception
-    mock_parent_del = MagicMock(side_effect=RuntimeError('Cleanup failed'))
-
-    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
-        # Trigger __del__ - should catch exception and not propagate it (line 104-106)
-        try:
-            tracking.__del__()
-            # Test passes if no exception is raised
-        except Exception as e:
-            raise AssertionError(
-                f'__del__ should handle exceptions gracefully, but raised: {e}'
-            ) from e
-
-
-@patch('model_manager.activities.experiment_tracking.Postgres.__init__')
-def test___del___handles_attribute_error_from_parent_del(mock_postgres_init):
-    """Test __del__ handles AttributeError from parent destructor - covers line 104 exception handling."""
-    mock_postgres_init.return_value = None
-
-    tracking = ExperimentTracking(
-        host='localhost',
-        port=5432,
-        user='test',
-        password='test',
-        dbname='test',
-        min_connections=1,
-        max_connections=10,
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-    )
-
-    # Add engine attribute
-    tracking.engine = MagicMock()
-
-    # Mock parent __del__ to raise AttributeError
-    mock_parent_del = MagicMock(side_effect=AttributeError('engine not found'))
-
-    with patch.object(type(tracking).__bases__[0], '__del__', mock_parent_del, create=True):
-        # Trigger __del__ - should catch AttributeError and not propagate it
-        try:
-            tracking.__del__()
-            # Test passes if no exception is raised
-        except Exception as e:
-            raise AssertionError(
-                f'__del__ should handle AttributeError gracefully, but raised: {e}'
-            ) from e

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 +0,0 @@
-"""Unit tests for Training activity."""
-
-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
-@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
-
-    # 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=[],
-    )
-
-    input_data = {
-        'metadata': {'workflow_id': 'test-123'},
-        'uploaded_file': uploaded_file,
-        'train_params': train_params,
-    }
-
-    # Execute
-    result = await training.train_model(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()
-
-
-@mark.asyncio
-@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
-
-    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=[],
-    )
-
-    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()
-
-    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': [],
-    }
-
-    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)
-
-    notification_handler.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)
-
-    training.info = MagicMock()
-
-    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': [],
-    }
-
-    result = await 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'

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/utils/models/test_train_model_params.py

@@ -1,497 +0,0 @@
-"""Unit tests for TrainModelParams class."""
-
-import pytest
-
-from model_manager.utils.models.train_model_params import TrainModelParams
-
-
-@pytest.fixture
-def valid_params_dict():
-    """Create valid parameters dictionary for testing."""
-    return {
-        'variable_columns': ['var1', 'var2', 'var3'],
-        '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},
-        'window': 10,
-        'use_scaler': True,
-        'include_ar': False,
-        'bucket_name': 'test-bucket',
-        'file_name': 'test-file.csv',
-        'line_separator': '\n',
-        'decimal_separator': '.',
-        'train_size': 80,
-        'shuffle': True,
-        'experiment_run_id': 123,
-        'experiment_name': 'Test experiment name',
-        '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)
-
-    assert params.variable_columns == ['var1', 'var2', 'var3']
-    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.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.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.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)
-
-    assert params.variable_columns == ['var1', 'var2', 'var3']
-    assert params.lag_train == 5
-    assert params.experiment_run_id == 123
-
-
-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
-
-    with pytest.raises(ValueError, match='variable_columns is required and cannot be None'):
-        TrainModelParams.from_dict(valid_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'
-
-    with pytest.raises(TypeError, match='variable_columns must be of type list'):
-        TrainModelParams.from_dict(valid_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
-
-    with pytest.raises(ValueError, match='lag_train is required and cannot be None'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-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)
-    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'
-
-    with pytest.raises(TypeError, match='removed_intervals must be of type list'):
-        TrainModelParams.from_dict(valid_params_dict)
-
-
-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'),
-    ]
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-    assert len(params.removed_intervals) == 2
-    assert params.removed_intervals[0] == ('2023-01-01', '2023-01-10')
-
-
-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'):
-        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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='train_size must be between 1 and 99'):
-        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'] = []
-
-    params = TrainModelParams.from_dict(valid_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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='lag_train must be positive'):
-        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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='lag_train must be positive'):
-        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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    with pytest.raises(ValueError, match='lag_val must be positive'):
-        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'
-
-    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)
-
-    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'
-
-    params = TrainModelParams.from_dict(valid_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'
-
-    params = TrainModelParams.from_dict(valid_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'
-
-    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)
-
-    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'
-
-    params = TrainModelParams.from_dict(valid_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'
-
-    params = TrainModelParams.from_dict(valid_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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    # Should not raise any exception
-    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'
-
-    params = TrainModelParams.from_dict(valid_params_dict)
-
-    # Should not raise any exception
-    params.validate_business_rules()

tests/utils/repository/test_model_repository.py

@@ -1,708 +0,0 @@
-from unittest.mock import MagicMock, patch
-
-import numpy as np
-import pytest
-from pandas import DataFrame
-
-from model_manager.utils.repository.model_repository import MLFlowRepository
-
-
-@pytest.fixture
-def mlflow_repository():
-    with patch(
-        'model_manager.utils.repository.model_repository.ModelServing', autospec=True
-    ) as mock_model_serving:
-        mock_instance = mock_model_serving.return_value
-        mock_instance.get_transformed_data = MagicMock()
-
-        repo = MLFlowRepository(
-            host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
-        )
-        return repo
-
-
-metadata = {
-    'metadata': {
-        'model_id': 'test_model',
-        'model_name': 'test_model',
-        'workflow_name': 'test_workflow',
-        'schema_name': 'test_schedule',
-    },
-}
-
-
-# ========== Tests for Model Artifact Generation Methods ==========
-
-
-def test_get_next_run_name(mlflow_repository):
-    """Test get_next_run_name generates correct run name based on existing runs."""
-    mlflow_repository.model_serving.search_runs_by_name.return_value = [
-        MagicMock(),
-        MagicMock(),
-        MagicMock(),
-    ]
-
-    result = mlflow_repository.get_next_run_name('test_experiment')
-
-    mlflow_repository.model_serving.search_runs_by_name.assert_called_once_with(
-        experiment_names=['test_experiment'], order_by=['start_time desc']
-    )
-    assert result == 'test_experiment-4'
-
-
-def test_get_next_run_name_first_run(mlflow_repository):
-    """Test get_next_run_name for first run (no existing runs)."""
-    mlflow_repository.model_serving.search_runs_by_name.return_value = []
-
-    result = mlflow_repository.get_next_run_name('test_experiment')
-
-    assert result == 'test_experiment-1'
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_artifacts_success(mock_path, mlflow_repository):
-    """Test generate_artifacts successfully creates all artifacts."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    # Mock data
-    params = MagicMock(spec=TrainModelParams)
-    params.target_variable = 'target'
-    params.variable_columns = ['feat1', 'feat2']
-    params.experiment_name = 'test_exp'
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = 'test_run-1'
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1, 2], 'feat2': [3, 4]})
-    data.y_train = DataFrame({'target': [5, 6]})
-    data.x_test = DataFrame({'feat1': [7, 8], 'feat2': [9, 10]})
-    data.y_test = DataFrame({'target': [11, 12]})
-    data.regr = MagicMock()
-    data.regr.predict = MagicMock(return_value=np.array([5.1, 6.1]))
-    data.y_pred = np.array([11.1, 12.1])
-
-    # Mock path operations
-    mock_path.exists.return_value = True
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-
-    # Mock private methods
-    mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
-    mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
-    mlflow_repository._setup_run_directory = MagicMock()
-    mlflow_repository._generate_report = MagicMock(return_value=data)
-
-    result = mlflow_repository.generate_artifacts(data)
-
-    # Assertions
-    mlflow_repository._get_reports_directory.assert_called_once()
-    mlflow_repository._create_run_directory.assert_called_once_with('/reports', 'test_run-1')
-    mlflow_repository._setup_run_directory.assert_called_once()
-    mlflow_repository._generate_report.assert_called_once()
-    assert result == data
-
-
-def test_generate_artifacts_missing_run_name(mlflow_repository):
-    """Test generate_artifacts raises ValueError when run_name is not set."""
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = None
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository.generate_artifacts(data)
-
-    assert 'run_name must be set before generating artifacts' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_artifacts_reports_directory_not_exists(mock_path, mlflow_repository):
-    """Test generate_artifacts raises FileNotFoundError when reports directory doesn't exist."""
-    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.target_variable = 'target'
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = 'test_run-1'
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame({'target': [2]})
-    data.x_test = DataFrame({'feat1': [3]})
-    data.y_test = DataFrame({'target': [4]})
-    data.regr = MagicMock()
-    data.y_pred = np.array([4.1])
-
-    mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
-    mock_path.exists.return_value = False
-
-    with pytest.raises(FileNotFoundError) as exc_info:
-        mlflow_repository.generate_artifacts(data)
-
-    assert 'Reports directory does not exist' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_artifacts_header_file_not_exists(mock_path, mlflow_repository):
-    """Test generate_artifacts raises FileNotFoundError when header.html doesn't exist."""
-    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.target_variable = 'target'
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = 'test_run-1'
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame({'target': [2]})
-    data.x_test = DataFrame({'feat1': [3]})
-    data.y_test = DataFrame({'target': [4]})
-    data.regr = MagicMock()
-    data.regr.predict = MagicMock(return_value=np.array([2.1]))
-    data.y_pred = np.array([4.1])
-
-    mlflow_repository._get_reports_directory = MagicMock(return_value='/reports')
-    mlflow_repository._create_run_directory = MagicMock(return_value='/reports/test_run-1_20231010')
-
-    # First call returns True (reports dir exists), second returns False (header.html doesn't exist)
-    mock_path.exists.side_effect = [True, False]
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-
-    with pytest.raises(FileNotFoundError) as exc_info:
-        mlflow_repository.generate_artifacts(data)
-
-    assert 'Header file does not exist' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_success(mock_path, mlflow_repository):
-    """Test save_run successfully logs all parameters, metrics, models, and 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.train_size = 80
-    params.removed_intervals = [(1, 10), (20, 30)]
-    params.experiment_name = 'test_exp'
-    params.target_variable = 'target'
-    params.variable_columns = ['feat1', 'feat2']
-    params.lag_train = 5
-    params.lag_val = 3
-    params.window = 10
-    params.low_lim = 0.0
-    params.upp_lim = 1.0
-    params.include_ar = True
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = 'test_run-1'
-    data.params = params
-    data.report_path = '/reports/report.html'
-    data.train_data_path = '/reports/train.csv'
-    data.test_data_path = '/reports/test.csv'
-    data.mse_val = 0.123
-    data.r2_val = 0.987
-    data.mae_val = 0.456
-    data.scaler_dict = {'scaler': 'minmax'}
-    data.process_data = MagicMock()
-    data.regr = MagicMock()
-
-    mock_path.exists.return_value = True
-
-    mlflow_repository.save_run(data)
-
-    # Verify experiment was set
-    mlflow_repository.model_serving.set_experiment.assert_called_once_with('test_exp')
-
-    # Verify parameters were logged
-    assert mlflow_repository.model_serving.log_param.call_count == 13
-
-    # Verify metrics were logged
-    mlflow_repository.model_serving.log_metric.assert_any_call('MSE', 0.123)
-    mlflow_repository.model_serving.log_metric.assert_any_call('R2', 0.987)
-    mlflow_repository.model_serving.log_metric.assert_any_call('MAE', 0.456)
-
-    # Verify models were logged
-    mlflow_repository.model_serving.log_model.assert_any_call(data.process_data, 'data_model')
-    mlflow_repository.model_serving.log_model.assert_any_call(data.regr, 'prediction_model')
-
-    # Verify artifacts were logged
-    mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/report.html')
-    mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/train.csv')
-    mlflow_repository.model_serving.log_artifact.assert_any_call('/reports/test.csv')
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_missing_report_path(mock_path, mlflow_repository):
-    """Test save_run raises ValueError when report_path is missing."""
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    data = MagicMock(spec=TrainModelResult)
-    data.report_path = None
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository.save_run(data)
-
-    assert 'Report file does not exist' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_missing_metrics(mock_path, mlflow_repository):
-    """Test save_run raises ValueError when metrics are None."""
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    data = MagicMock(spec=TrainModelResult)
-    data.report_path = '/reports/report.html'
-    data.train_data_path = '/reports/train.csv'
-    data.test_data_path = '/reports/test.csv'
-    data.mse_val = None
-    data.r2_val = 0.987
-    data.mae_val = 0.456
-
-    mock_path.exists.return_value = True
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository.save_run(data)
-
-    assert 'One or more metrics (MSE, R2, MAE) are None' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_mlflow_error(mock_path, mlflow_repository):
-    """Test save_run handles MLflow errors 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.train_size = 80
-    params.removed_intervals = []
-    params.experiment_name = 'test_exp'
-
-    data = MagicMock(spec=TrainModelResult)
-    data.run_name = 'test_run-1'
-    data.params = params
-    data.report_path = '/reports/report.html'
-    data.train_data_path = '/reports/train.csv'
-    data.test_data_path = '/reports/test.csv'
-    data.mse_val = 0.123
-    data.r2_val = 0.987
-    data.mae_val = 0.456
-
-    mock_path.exists.return_value = True
-    mlflow_repository.model_serving.set_experiment.side_effect = Exception(
-        'MLflow connection error'
-    )
-
-    with pytest.raises(RuntimeError) as exc_info:
-        mlflow_repository.save_run(data)
-
-    assert 'Failed to save run' in str(exc_info.value)
-    assert 'MLflow connection error' in str(exc_info.value)
-
-
-# ========== Additional Tests for 100% Coverage ==========
-
-
-def test_init_artifacts_data_empty_x_train(mlflow_repository):
-    """Test _init_artifacts_data raises ValueError when x_train is empty."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame()  # Empty DataFrame
-    data.y_train = DataFrame({'target': [1]})
-    data.x_test = DataFrame({'feat1': [1]})
-    data.y_test = DataFrame({'target': [1]})
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._init_artifacts_data(data)
-
-    assert 'Training features (x_train) are empty' in str(exc_info.value)
-
-
-def test_init_artifacts_data_empty_y_train(mlflow_repository):
-    """Test _init_artifacts_data raises ValueError when y_train is empty."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame()  # Empty DataFrame
-    data.x_test = DataFrame({'feat1': [1]})
-    data.y_test = DataFrame({'target': [1]})
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._init_artifacts_data(data)
-
-    assert 'Training target (y_train) is empty' in str(exc_info.value)
-
-
-def test_init_artifacts_data_empty_x_test(mlflow_repository):
-    """Test _init_artifacts_data raises ValueError when x_test is empty."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame({'target': [1]})
-    data.x_test = DataFrame()  # Empty DataFrame
-    data.y_test = DataFrame({'target': [1]})
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._init_artifacts_data(data)
-
-    assert 'Test features (x_test) are empty' in str(exc_info.value)
-
-
-def test_init_artifacts_data_empty_y_test(mlflow_repository):
-    """Test _init_artifacts_data raises ValueError when y_test is empty."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame({'target': [1]})
-    data.x_test = DataFrame({'feat1': [1]})
-    data.y_test = DataFrame()  # Empty DataFrame
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._init_artifacts_data(data)
-
-    assert 'Test target (y_test) is empty' in str(exc_info.value)
-
-
-def test_init_artifacts_data_none_y_pred(mlflow_repository):
-    """Test _init_artifacts_data raises ValueError when y_pred is None."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1]})
-    data.y_train = DataFrame({'target': [1]})
-    data.x_test = DataFrame({'feat1': [1]})
-    data.y_test = DataFrame({'target': [1]})
-    data.y_pred = None
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._init_artifacts_data(data)
-
-    assert 'Test predictions (y_pred) are None' in str(exc_info.value)
-
-
-def test_init_artifacts_data_success(mlflow_repository):
-    """Test _init_artifacts_data successfully prepares data."""
-    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.target_variable = 'target'
-
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.x_train = DataFrame({'feat1': [1, 2]})
-    data.y_train = DataFrame({'target': [3, 4]})
-    data.x_test = DataFrame({'feat1': [5, 6]})
-    data.y_test = DataFrame({'target': [7, 8]})
-    data.regr = MagicMock()
-    data.regr.predict = MagicMock(return_value=np.array([3.1, 4.1]))
-    data.y_pred = np.array([7.1, 8.1])
-
-    reference_data, current_data = mlflow_repository._init_artifacts_data(data)
-
-    assert 'target' in reference_data.columns
-    assert 'prediction' in reference_data.columns
-    assert 'target' in current_data.columns
-    assert 'prediction' in current_data.columns
-    assert len(reference_data) == 2
-    assert len(current_data) == 2
-
-
-@patch('model_manager.utils.repository.model_repository.makedirs')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_create_run_directory_success(mock_path, mock_makedirs, mlflow_repository):
-    """Test _create_run_directory successfully creates directory."""
-    mock_path.join.return_value = '/reports/test_run_20231010_123456_123456'
-
-    result = mlflow_repository._create_run_directory('/reports', 'test_run')
-
-    mock_makedirs.assert_called_once_with('/reports/test_run_20231010_123456_123456', exist_ok=True)
-    assert result == '/reports/test_run_20231010_123456_123456'
-
-
-@patch('model_manager.utils.repository.model_repository.makedirs')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_create_run_directory_permission_error(mock_path, mock_makedirs, mlflow_repository):
-    """Test _create_run_directory handles PermissionError."""
-    mock_path.join.return_value = '/reports/test_run_20231010'
-    mock_makedirs.side_effect = PermissionError('Permission denied')
-
-    with pytest.raises(PermissionError) as exc_info:
-        mlflow_repository._create_run_directory('/reports', 'test_run')
-
-    assert 'Permission denied when creating directory' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.makedirs')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_create_run_directory_os_error(mock_path, mock_makedirs, mlflow_repository):
-    """Test _create_run_directory handles OSError."""
-    mock_path.join.return_value = '/reports/test_run_20231010'
-    mock_makedirs.side_effect = OSError('Disk full')
-
-    with pytest.raises(OSError) as exc_info:
-        mlflow_repository._create_run_directory('/reports', 'test_run')
-
-    assert 'Failed to create directory' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.shutil')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_setup_run_directory_success(mock_path, mock_shutil, mlflow_repository):
-    """Test _setup_run_directory successfully sets up directory."""
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-    mock_open = MagicMock()
-
-    with patch('builtins.open', mock_open):
-        mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
-
-    assert mock_open.call_count == 3  # 3 empty files
-    mock_shutil.copy.assert_called_once_with('/reports/header.html', '/run_dir/header.html')
-
-
-@patch('model_manager.utils.repository.model_repository.shutil')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_setup_run_directory_file_not_found(mock_path, mock_shutil, mlflow_repository):
-    """Test _setup_run_directory handles FileNotFoundError."""
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-    mock_shutil.copy.side_effect = FileNotFoundError('Header not found')
-
-    mock_open = MagicMock()
-    with patch('builtins.open', mock_open):
-        with pytest.raises(FileNotFoundError) as exc_info:
-            mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
-
-    assert 'Header file not found' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.shutil')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_setup_run_directory_permission_error(mock_path, mock_shutil, mlflow_repository):
-    """Test _setup_run_directory handles PermissionError."""
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-
-    mock_open = MagicMock()
-    mock_open.side_effect = PermissionError('Permission denied')
-
-    with patch('builtins.open', mock_open):
-        with pytest.raises(PermissionError) as exc_info:
-            mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
-
-    assert 'Permission denied when setting up directory' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.shutil')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_setup_run_directory_os_error(mock_path, mock_shutil, mlflow_repository):
-    """Test _setup_run_directory handles OSError."""
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-
-    mock_open = MagicMock()
-    mock_open.side_effect = OSError('Disk error')
-
-    with patch('builtins.open', mock_open):
-        with pytest.raises(OSError) as exc_info:
-            mlflow_repository._setup_run_directory('/run_dir', '/reports/header.html')
-
-    assert 'Failed to setup run directory' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.Reports')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_report_success(mock_path, mock_reports, mlflow_repository):
-    """Test _generate_report successfully generates all reports."""
-    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.variable_columns = ['feat1', 'feat2']
-
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.run_dir = '/run_dir'
-
-    reference_data = DataFrame(
-        {'feat1': [1.0], 'feat2': [2.0], 'target': [3.0], 'prediction': [3.1]}
-    )
-    current_data = DataFrame({'feat1': [4.0], 'feat2': [5.0], 'target': [6.0], 'prediction': [6.1]})
-
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-    mock_report_instance = MagicMock()
-    mock_reports.return_value = mock_report_instance
-
-    # Mock DataFrame.to_csv to avoid actual file writing
-    with patch.object(DataFrame, 'to_csv'):
-        result = mlflow_repository._generate_report(reference_data, current_data, data)
-
-    mock_reports.assert_called_once()
-    mock_report_instance.add_data_quality_section.assert_called_once()
-    mock_report_instance.add_data_drift_section.assert_called_once()
-    mock_report_instance.add_regression_section.assert_called_once()
-    mock_report_instance.save_all_sections_html.assert_called_once()
-    assert result == data
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_report_value_error(mock_path, mlflow_repository):
-    """Test _generate_report handles ValueError from data conversion."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    params = MagicMock(spec=TrainModelParams)
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.run_dir = '/run_dir'
-
-    # DataFrame with non-numeric data
-    reference_data = DataFrame({'feat1': ['a', 'b']})
-    current_data = DataFrame({'feat1': ['c', 'd']})
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._generate_report(reference_data, current_data, data)
-
-    assert 'Failed to convert data to float64' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.Reports')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_report_permission_error(mock_path, mock_reports, mlflow_repository):
-    """Test _generate_report handles PermissionError."""
-    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.variable_columns = ['feat1']
-
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.run_dir = '/run_dir'
-
-    reference_data = DataFrame({'feat1': [1.0]})
-    current_data = DataFrame({'feat1': [2.0]})
-
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-    mock_report_instance = MagicMock()
-    mock_reports.return_value = mock_report_instance
-    mock_report_instance.save_all_sections_html.side_effect = PermissionError('Permission denied')
-
-    with pytest.raises(PermissionError) as exc_info:
-        mlflow_repository._generate_report(reference_data, current_data, data)
-
-    assert 'Permission denied when writing report files' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.Reports')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_report_os_error(mock_path, mock_reports, mlflow_repository):
-    """Test _generate_report handles OSError."""
-    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.variable_columns = ['feat1']
-
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.run_dir = '/run_dir'
-
-    reference_data = DataFrame({'feat1': [1.0]})
-    current_data = DataFrame({'feat1': [2.0]})
-
-    mock_path.join.side_effect = lambda *args: '/'.join(args)
-    mock_report_instance = MagicMock()
-    mock_reports.return_value = mock_report_instance
-    mock_report_instance.save_all_sections_html.side_effect = OSError('Disk error')
-
-    with pytest.raises(OSError) as exc_info:
-        mlflow_repository._generate_report(reference_data, current_data, data)
-
-    assert 'Failed to generate report' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.Reports')
-@patch('model_manager.utils.repository.model_repository.path')
-def test_generate_report_run_dir_none(mock_path, mock_reports, mlflow_repository):
-    """Test _generate_report raises ValueError when run_dir is None."""
-    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.variable_columns = ['feat1']
-
-    data = MagicMock(spec=TrainModelResult)
-    data.params = params
-    data.run_dir = None  # Not set
-
-    reference_data = DataFrame({'feat1': [1.0]})
-    current_data = DataFrame({'feat1': [2.0]})
-
-    mock_report_instance = MagicMock()
-    mock_reports.return_value = mock_report_instance
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository._generate_report(reference_data, current_data, data)
-
-    assert 'run_dir is not set after directory creation' in str(exc_info.value)
-
-
-def test_get_reports_directory(mlflow_repository):
-    """Test _get_reports_directory returns correct path."""
-    result = mlflow_repository._get_reports_directory()
-
-    assert result.endswith('model_manager/reports')
-    assert 'model_manager' in result
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_missing_train_data_path(mock_path, mlflow_repository):
-    """Test save_run raises ValueError when train_data_path is missing."""
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    data = MagicMock(spec=TrainModelResult)
-    data.report_path = '/reports/report.html'
-    data.train_data_path = None
-
-    mock_path.exists.return_value = True
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository.save_run(data)
-
-    assert 'Training data file does not exist' in str(exc_info.value)
-
-
-@patch('model_manager.utils.repository.model_repository.path')
-def test_save_run_missing_test_data_path(mock_path, mlflow_repository):
-    """Test save_run raises ValueError when test_data_path is missing."""
-    from model_manager.utils.models.train_model_result import TrainModelResult
-
-    data = MagicMock(spec=TrainModelResult)
-    data.report_path = '/reports/report.html'
-    data.train_data_path = '/reports/train.csv'
-    data.test_data_path = None
-
-    mock_path.exists.return_value = True
-
-    with pytest.raises(ValueError) as exc_info:
-        mlflow_repository.save_run(data)
-
-    assert 'Test data file does not exist' in str(exc_info.value)

tests/utils/repository/test_training_repository.py

@@ -1,324 +0,0 @@
-"""Unit tests for TrainingRepository."""
-
-from io import BytesIO
-from unittest.mock import MagicMock, patch
-
-import numpy as np
-import pandas as pd
-from pytest import fixture, raises
-
-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.training_repository import TrainingRepository
-
-
-@fixture
-def logger():
-    """Create a mock logger."""
-    return MagicMock()
-
-
-@fixture
-def training_repository(logger):
-    """Create a TrainingRepository instance."""
-    return TrainingRepository(logger)
-
-
-@fixture
-def train_params():
-    """Create sample training parameters."""
-    return TrainModelParams(
-        variable_columns=['feature1', 'feature2'],
-        lag_train=1,
-        lag_val=1,
-        target_variable='target',
-        rem_static_win=False,
-        low_lim={'feature1': 0.0, 'feature2': 0.0},
-        upp_lim={'feature1': 100.0, 'feature2': 100.0},
-        window=10,
-        use_scaler=True,
-        include_ar=False,
-        bucket_name='test-bucket',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        train_size=80,
-        shuffle=True,
-        experiment_run_id=123,
-        experiment_name='test_experiment',
-        removed_intervals=[],
-    )
-
-
-@fixture
-def sample_csv_data():
-    """Create sample CSV data."""
-    csv_content = """feature1,feature2,target
-1.0,2.0,10.0
-2.0,3.0,15.0
-3.0,4.0,20.0
-4.0,5.0,25.0
-5.0,6.0,30.0
-6.0,7.0,35.0
-7.0,8.0,40.0
-8.0,9.0,45.0
-9.0,10.0,50.0
-10.0,11.0,55.0
-"""
-    return BytesIO(csv_content.encode())
-
-
-@patch('model_manager.utils.repository.training_repository.load_data')
-@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
-@patch('model_manager.utils.repository.training_repository.split_train_test')
-@patch('model_manager.utils.repository.training_repository.LinearRegressionModel')
-def test_train_success(
-    mock_linear_model,
-    mock_split,
-    mock_preprocessor_class,
-    mock_load_data,
-    training_repository,
-    train_params,
-    sample_csv_data,
-):
-    """Test successful model training."""
-    # Setup mocks
-    mock_data = pd.DataFrame(
-        {'feature1': [1, 2, 3, 4, 5], 'feature2': [2, 3, 4, 5, 6], 'target': [10, 15, 20, 25, 30]}
-    )
-    mock_load_data.return_value = mock_data
-
-    mock_preprocessor = MagicMock()
-    mock_preprocessor_class.return_value = mock_preprocessor
-    mock_preprocessor.transform.return_value = mock_data
-
-    x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
-    x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
-    y_train = pd.Series([10, 15, 20], name='target')
-    y_test = pd.Series([25, 30], name='target')
-    mock_split.return_value = (x_train, x_test, y_train, y_test)
-
-    mock_model = MagicMock()
-    mock_linear_model.return_value = mock_model
-
-    mock_scaler = MagicMock()
-    mock_preprocessor.get_scaler.return_value = mock_scaler
-
-    # Execute
-    result = training_repository.train(sample_csv_data, train_params)
-
-    # Assertions
-    assert isinstance(result, TrainModelResult)
-    assert result.params == train_params
-    assert result.process_data == mock_preprocessor
-    assert result.regr == mock_model
-    mock_load_data.assert_called_once()
-    mock_preprocessor.fit.assert_called_once()
-    mock_model.fit.assert_called_once()
-
-
-@patch('model_manager.utils.repository.training_repository.load_data')
-def test_train_empty_data_after_transform(
-    mock_load_data, training_repository, train_params, sample_csv_data
-):
-    """Test training with empty data after transformation."""
-    mock_data = pd.DataFrame({'feature1': [], 'feature2': [], 'target': []})
-    mock_load_data.return_value = mock_data
-
-    with patch.object(training_repository, 'init_data_preprocessor') as mock_init:
-        mock_preprocessor = MagicMock()
-        mock_init.return_value = mock_preprocessor
-        mock_preprocessor.transform.return_value = pd.DataFrame()
-
-        with raises(ValueError, match='Data view is empty after transformation'):
-            training_repository.train(sample_csv_data, train_params)
-
-
-def test_init_scaler_dict_with_minmax_scaler(training_repository, train_params):
-    """Test scaler dict initialization with MinMaxScaler."""
-    mock_preprocessor = MagicMock()
-    mock_scaler = MagicMock()
-    mock_scaler.x_min = [0.0, 1.0]
-    mock_scaler.x_max = [10.0, 11.0]
-    mock_scaler.y_min = 5.0
-    mock_scaler.y_max = 50.0
-    mock_preprocessor.get_scaler.return_value = mock_scaler
-
-    # Patch isinstance to return True for MinMaxScaler
-    with patch(
-        'model_manager.utils.repository.training_repository.isinstance',
-        side_effect=lambda obj, cls: cls.__name__ == 'MinMaxScaler',
-    ):
-        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
-
-    assert result is not None
-    assert 'feature1' in result
-    assert 'feature2' in result
-    assert 'target' in result
-    assert result['feature1'] == {'min': 0.0, 'max': 10.0}
-    assert result['feature2'] == {'min': 1.0, 'max': 11.0}
-    assert result['target'] == {'min': 5.0, 'max': 50.0}
-
-
-def test_init_scaler_dict_with_z_scaler(training_repository, train_params):
-    """Test scaler dict initialization with Z_Scaler."""
-    mock_preprocessor = MagicMock()
-    mock_scaler = MagicMock()
-    mock_scaler.create_dict.return_value = {'mean': 5.0, 'std': 2.0}
-    mock_preprocessor.get_scaler.return_value = mock_scaler
-
-    # Patch isinstance to return True for Z_Scaler
-    with patch(
-        'model_manager.utils.repository.training_repository.isinstance',
-        side_effect=lambda obj, cls: cls.__name__ == 'Z_Scaler',
-    ):
-        result = training_repository.init_scaler_dict(mock_preprocessor, train_params)
-
-    assert result == {'mean': 5.0, 'std': 2.0}
-    mock_scaler.create_dict.assert_called_once()
-
-
-def test_init_scaler_dict_without_scaler(training_repository):
-    """Test scaler dict initialization when use_scaler is False."""
-    train_params_no_scaler = TrainModelParams(
-        variable_columns=['feature1'],
-        lag_train=1,
-        lag_val=1,
-        target_variable='target',
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        train_size=80,
-        shuffle=True,
-        experiment_run_id=123,
-        experiment_name='test',
-        removed_intervals=[],
-    )
-
-    mock_preprocessor = MagicMock()
-    result = training_repository.init_scaler_dict(mock_preprocessor, train_params_no_scaler)
-
-    assert result == {}
-
-
-@patch('model_manager.utils.repository.training_repository.mse')
-@patch('model_manager.utils.repository.training_repository.mae')
-@patch('model_manager.utils.repository.training_repository.r2')
-def test_after_train_calculation_with_scaler(
-    mock_r2, mock_mae, mock_mse, training_repository, train_params
-):
-    """Test post-training calculations with scaler."""
-    # Setup mock train result
-    mock_train_result = MagicMock(spec=TrainModelResult)
-    mock_train_result.params = train_params
-    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3], 'feature2': [2, 3, 4]})
-    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5], 'feature2': [5, 6]})
-    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
-    mock_train_result.y_test = pd.Series([25, 30], name='target')
-
-    mock_regr = MagicMock()
-    mock_regr.predict.return_value = np.array([24.5, 29.5])
-    mock_train_result.regr = mock_regr
-
-    mock_scaler = MagicMock()
-    mock_scaler.denormalize_single_input.side_effect = lambda x, col: x
-    mock_scaler.denormalize_predictions.side_effect = lambda x, col: x
-
-    mock_process_data = MagicMock()
-    mock_process_data.get_scaler.return_value = mock_scaler
-    mock_train_result.process_data = mock_process_data
-
-    # Setup metric mocks
-    mock_mse.return_value = 0.5
-    mock_mae.return_value = 0.3
-    mock_r2.return_value = 0.95
-
-    # Execute
-    result = training_repository.after_train_calculation(train_params, mock_train_result)
-
-    # Assertions
-    assert result == mock_train_result
-    assert result.mse_val == 0.5
-    assert result.mae_val == 0.3
-    assert result.r2_val == 0.95
-    assert result.y_pred is not None
-    mock_regr.predict.assert_called_once()
-    mock_mse.assert_called_once()
-    mock_mae.assert_called_once()
-    mock_r2.assert_called_once()
-
-
-@patch('model_manager.utils.repository.training_repository.mse')
-@patch('model_manager.utils.repository.training_repository.mae')
-@patch('model_manager.utils.repository.training_repository.r2')
-def test_after_train_calculation_without_scaler(mock_r2, mock_mae, mock_mse, training_repository):
-    """Test post-training calculations without scaler."""
-    train_params_no_scaler = TrainModelParams(
-        variable_columns=['feature1'],
-        lag_train=1,
-        lag_val=1,
-        target_variable='target',
-        rem_static_win=False,
-        low_lim={'feature1': 0.0},
-        upp_lim={'feature1': 100.0},
-        window=10,
-        use_scaler=False,
-        include_ar=False,
-        bucket_name='test',
-        file_name='test.csv',
-        line_separator='\n',
-        decimal_separator='.',
-        train_size=80,
-        shuffle=True,
-        experiment_run_id=123,
-        experiment_name='test',
-        removed_intervals=[],
-    )
-
-    mock_train_result = MagicMock(spec=TrainModelResult)
-    mock_train_result.params = train_params_no_scaler
-    mock_train_result.x_train = pd.DataFrame({'feature1': [1, 2, 3]})
-    mock_train_result.x_test = pd.DataFrame({'feature1': [4, 5]})
-    mock_train_result.y_train = pd.Series([10, 15, 20], name='target')
-    mock_train_result.y_test = pd.Series([25, 30], name='target')
-
-    mock_regr = MagicMock()
-    mock_regr.predict.return_value = np.array([24.5, 29.5])
-    mock_train_result.regr = mock_regr
-
-    # Setup metric mocks
-    mock_mse.return_value = 0.5
-    mock_mae.return_value = 0.3
-    mock_r2.return_value = 0.95
-
-    # Execute
-    result = training_repository.after_train_calculation(train_params_no_scaler, mock_train_result)
-
-    # Assertions
-    assert result.mse_val == 0.5
-    assert result.mae_val == 0.3
-    assert result.r2_val == 0.95
-
-
-@patch('model_manager.utils.repository.training_repository.DataPreprocessor')
-def test_init_data_preprocessor(mock_preprocessor_class, training_repository, train_params):
-    """Test DataPreprocessor initialization."""
-    mock_preprocessor = MagicMock()
-    mock_preprocessor_class.return_value = mock_preprocessor
-
-    result = training_repository.init_data_preprocessor(train_params)
-
-    assert result == mock_preprocessor
-    mock_preprocessor_class.assert_called_once()
-    call_kwargs = mock_preprocessor_class.call_args[1]
-    assert call_kwargs['target_variable'] == 'target'
-    assert call_kwargs['input_columns'] == ['feature1', 'feature2']
-    assert call_kwargs['low_lim'] == {'feature1': 0.0, 'feature2': 0.0}
-    assert call_kwargs['upp_lim'] == {'feature1': 100.0, 'feature2': 100.0}

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/worker/test_worker.py

@@ -1,494 +0,0 @@
-from unittest.mock import ANY, AsyncMock, MagicMock, 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')
-
-
-@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
-    start_prometheus_server()
-
-    # Assert
-    mock_start_http_server.assert_called_once_with(9090)
-    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)
-
-
-@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
-    from model_manager.worker.worker import start_prometheus_server
-
-    # Act
-    start_prometheus_server()
-
-    # Assert
-    mock_start_http_server.assert_called_once_with(9090)
-    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.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_success(
-    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 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
-    from model_manager.worker.worker import main
-
-    # Act
-    await main()
-
-    # Assert
-    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)
-
-
-@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.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')
-@patch('model_manager.worker.worker.sys.exit')
-@patch('model_manager.worker.worker.metrics')
-async def test_main_exception_handling(
-    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_env_vars,
-):
-    """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
-    from model_manager.worker.worker import main
-
-    # Act
-    await main()
-
-    # Assert - Verify cleanup was performed
-    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)
-
-
-@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.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_creates_only_one_worker(
-    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,
-):
-    """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
-    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 only one worker was created
-    assert mock_worker.call_count == 1
-
-    # 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'])
-
-
-@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.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,
-):
-    """Test that main initializes Activities 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
-    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
-
-
-@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.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,
-):
-    """Test that main uses environment variables correctly."""
-    # 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
-    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)

tests/workflows/test_train_model.py

@@ -1,678 +0,0 @@
-"""Unit tests for TrainModel workflow."""
-
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-from pytest import fixture, mark
-
-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
-from model_manager.workflows.train_model import TrainModel
-
-
-@fixture
-def train_model_workflow() -> TrainModel:
-    """Fixture for TrainModel workflow instance."""
-    return TrainModel()
-
-
-@fixture
-def mock_train_params():
-    """Fixture for mock TrainModelParams."""
-    return 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=[],
-    )
-
-
-@fixture
-def mock_train_result(mock_train_params):
-    """Fixture for mock TrainModelResult."""
-    result = MagicMock(spec=TrainModelResult)
-    result.params = mock_train_params
-    result.run_name = 'test_experiment-1'
-    result.run_dir = 'test_run_dir'  # Relative path instead of /tmp
-    return result
-
-
-# ============================================================================
-# Tests for run() - Complete workflow
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_run_success_complete_flow(
-    workflow_mock: AsyncMock,
-    train_model_workflow: TrainModel,
-    mock_train_params,
-    mock_train_result,
-):
-    """Test successful complete workflow execution."""
-    input_data = {
-        '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': [],
-    }
-
-    # Mock workflow.logger to avoid RuntimeWarning about unawaited coroutines
-    workflow_mock.logger = MagicMock()
-
-    # Mock activity responses
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_train_params,  # validate_train_params
-            None,  # update_experiment_run (MAGE_WAITING_PROC)
-            b'file_content',  # fetch_file_from_minio
-            mock_train_result,  # train_model
-            None,  # update_experiment_run (TRAINING_SUCCESS)
-            mock_train_result,  # save_model
-            None,  # update_experiment_run (MLFLOW_SENT with run_name)
-            None,  # cleanup_run_directory
-            None,  # delete_file_from_minio
-            None,  # update_experiment_run (FILE_DELETED)
-        ]
-    )
-
-    # Execute workflow
-    await train_model_workflow.run(input_data)
-
-    # Verify all activity calls (now 10 instead of 9 due to cleanup_run_directory)
-    assert workflow_mock.execute_activity_method.call_count == 10
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_run_missing_experiment_run_id(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test workflow fails when experiment_run_id is missing."""
-    input_data = {
-        'target_variable': 'price',
-        'variable_columns': ['feature1'],
-    }
-
-    # Mock workflow.logger to avoid NotInWorkflowEventLoopError
-    workflow_mock.logger = MagicMock()
-
-    with pytest.raises(ValueError, match='experiment_run_id is required'):
-        await train_model_workflow.run(input_data)
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_run_invalid_experiment_run_id_type(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test workflow fails when experiment_run_id has invalid type."""
-    input_data = {
-        'experiment_run_id': 'invalid',  # Should be int
-        'target_variable': 'price',
-    }
-
-    # Mock workflow.logger to avoid NotInWorkflowEventLoopError
-    workflow_mock.logger = MagicMock()
-
-    with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
-        await train_model_workflow.run(input_data)
-
-
-# ============================================================================
-# Tests for _validate_experiment_run_id()
-# ============================================================================
-
-
-@patch('model_manager.workflows.train_model.workflow')
-def test_validate_experiment_run_id_success(
-    workflow_mock: MagicMock, train_model_workflow: TrainModel
-):
-    """Test successful experiment_run_id validation."""
-    workflow_mock.logger = MagicMock()
-    input_data = {'experiment_run_id': 456}
-
-    result = train_model_workflow._validate_experiment_run_id(input_data)
-
-    assert result == 456
-
-
-@patch('model_manager.workflows.train_model.workflow')
-def test_validate_experiment_run_id_missing(
-    workflow_mock: MagicMock, train_model_workflow: TrainModel
-):
-    """Test validation fails when experiment_run_id is missing."""
-    workflow_mock.logger = MagicMock()
-    input_data = {}
-
-    with pytest.raises(ValueError, match='experiment_run_id is required'):
-        train_model_workflow._validate_experiment_run_id(input_data)
-
-
-@patch('model_manager.workflows.train_model.workflow')
-def test_validate_experiment_run_id_none(
-    workflow_mock: MagicMock, train_model_workflow: TrainModel
-):
-    """Test validation fails when experiment_run_id is None."""
-    workflow_mock.logger = MagicMock()
-    input_data = {'experiment_run_id': None}
-
-    with pytest.raises(ValueError, match='experiment_run_id is required'):
-        train_model_workflow._validate_experiment_run_id(input_data)
-
-
-@patch('model_manager.workflows.train_model.workflow')
-def test_validate_experiment_run_id_invalid_type(
-    workflow_mock: MagicMock, train_model_workflow: TrainModel
-):
-    """Test validation fails when experiment_run_id is not an integer."""
-    workflow_mock.logger = MagicMock()
-    input_data = {'experiment_run_id': 'not_an_int'}
-
-    with pytest.raises(ValueError, match='experiment_run_id must be an integer'):
-        train_model_workflow._validate_experiment_run_id(input_data)
-
-
-# ============================================================================
-# Tests for _validate_training_parameters()
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_validate_training_parameters_success(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
-):
-    """Test successful parameter validation."""
-    input_data = {
-        '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',
-        'removed_intervals': [],
-    }
-
-    metadata = {
-        'metadata': {
-            'experiment_run_id': 123,
-            'workflow_name': 'train_model',
-        }
-    }
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_train_params,  # validate_train_params
-            None,  # update_experiment_run
-        ]
-    )
-
-    result = await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
-
-    assert result == mock_train_params
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_validate_training_parameters_validation_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test parameter validation handles errors correctly."""
-    input_data = {'experiment_run_id': 123}
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    # Mock workflow.logger to avoid RuntimeWarning about unawaited coroutines
-    workflow_mock.logger = MagicMock()
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            ValueError('Missing required field'),  # validate_train_params fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(ValueError, match='Missing required field'):
-        await train_model_workflow._validate_training_parameters(input_data, 123, metadata)
-
-    # Verify error status was updated
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-# ============================================================================
-# Tests for _download_and_train_model()
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_success(
-    workflow_mock: AsyncMock,
-    train_model_workflow: TrainModel,
-    mock_train_params,
-    mock_train_result,
-):
-    """Test successful download and training."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            b'file_content',  # fetch_file_from_minio
-            mock_train_result,  # train_model
-            None,  # update_experiment_run
-        ]
-    )
-
-    result = await train_model_workflow._download_and_train_model(
-        train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-    )
-
-    assert result == mock_train_result
-    assert workflow_mock.execute_activity_method.call_count == 3
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_download_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
-):
-    """Test download error is handled correctly."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    # Mock workflow.logger to avoid RuntimeWarning about unawaited coroutines
-    workflow_mock.logger = MagicMock()
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            Exception('MinIO connection failed'),  # fetch_file_from_minio fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(Exception, match='MinIO connection failed'):
-        await train_model_workflow._download_and_train_model(
-            train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-        )
-
-    # Verify error status was updated
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_training_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
-):
-    """Test training error is handled correctly."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    # Mock workflow.logger to avoid RuntimeWarning about unawaited coroutines
-    workflow_mock.logger = MagicMock()
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            b'file_content',  # fetch_file_from_minio succeeds
-            Exception('Training failed'),  # train_model fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(Exception, match='Training failed'):
-        await train_model_workflow._download_and_train_model(
-            train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-        )
-
-    # Verify error status was updated
-    assert workflow_mock.execute_activity_method.call_count == 3
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_closes_bytesio_on_success(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
-):
-    """Test that BytesIO is closed in finally block on success."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    # Create a mock BytesIO with close method
-    mock_file = MagicMock()
-    mock_file.close = MagicMock()
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_file,  # fetch_file_from_minio returns BytesIO
-            mock_train_result,  # train_model succeeds
-            None,  # update_experiment_run
-        ]
-    )
-
-    result = await train_model_workflow._download_and_train_model(
-        train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-    )
-
-    # Verify BytesIO.close() was called in finally block
-    mock_file.close.assert_called_once()
-    assert result == mock_train_result
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_closes_bytesio_on_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params
-):
-    """Test that BytesIO is closed in finally block even on error."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    # Create a mock BytesIO with close method
-    mock_file = MagicMock()
-    mock_file.close = MagicMock()
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_file,  # fetch_file_from_minio returns BytesIO
-            Exception('Training failed'),  # train_model fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(Exception, match='Training failed'):
-        await train_model_workflow._download_and_train_model(
-            train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-        )
-
-    # Verify BytesIO.close() was called in finally block even after exception
-    mock_file.close.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_download_and_train_model_handles_file_without_close(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result
-):
-    """Test that workflow handles file objects without close method gracefully."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    # Create a mock file without close method
-    mock_file = MagicMock(spec=[])
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_file,  # fetch_file_from_minio returns object without close
-            mock_train_result,  # train_model succeeds
-            None,  # update_experiment_run
-        ]
-    )
-
-    # Should not raise error even if file doesn't have close method
-    result = await train_model_workflow._download_and_train_model(
-        train_params=mock_train_params, experiment_run_id=123, metadata=metadata
-    )
-
-    assert result == mock_train_result
-
-
-# ============================================================================
-# Tests for _save_model_to_mlflow()
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_save_model_to_mlflow_success(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
-):
-    """Test successful model saving to MLFlow."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            mock_train_result,  # save_model
-            None,  # update_experiment_run with MODEL_SAVED
-        ]
-    )
-
-    result = await train_model_workflow._save_model_to_mlflow(
-        train_result=mock_train_result, experiment_run_id=123, metadata=metadata
-    )
-
-    assert result == mock_train_result
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_save_model_to_mlflow_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result
-):
-    """Test MLFlow save error is handled correctly."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            Exception('MLFlow connection failed'),  # save_model fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(Exception, match='MLFlow connection failed'):
-        await train_model_workflow._save_model_to_mlflow(
-            train_result=mock_train_result, experiment_run_id=123, metadata=metadata
-        )
-
-    # Verify error status was updated
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-# ============================================================================
-# Tests for _cleanup_resources()
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_cleanup_resources_success(
-    workflow_mock: AsyncMock,
-    train_model_workflow: TrainModel,
-    mock_train_result,
-):
-    """Test successful resource cleanup."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            None,  # cleanup_run_directory
-            None,  # delete_file_from_minio
-            None,  # update_experiment_run
-        ]
-    )
-
-    await train_model_workflow._cleanup_resources(
-        saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
-    )
-
-    # Verify activities were called (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
-    assert workflow_mock.execute_activity_method.call_count == 3
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_cleanup_resources_delete_error(
-    workflow_mock: AsyncMock,
-    train_model_workflow: TrainModel,
-    mock_train_result,
-):
-    """Test cleanup handles delete errors correctly."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    workflow_mock.logger = MagicMock()
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            None,  # cleanup_run_directory succeeds
-            Exception('MinIO delete failed'),  # delete_file_from_minio fails
-            None,  # update_experiment_run with error
-        ]
-    )
-
-    with pytest.raises(Exception, match='MinIO delete failed'):
-        await train_model_workflow._cleanup_resources(
-            saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
-        )
-
-    # Verify error status was updated (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
-    assert workflow_mock.execute_activity_method.call_count == 3
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_cleanup_resources_without_run_dir(
-    workflow_mock: AsyncMock,
-    train_model_workflow: TrainModel,
-    mock_train_result,
-):
-    """Test cleanup works when run_dir is not set."""
-    metadata = {'metadata': {'experiment_run_id': 123, 'workflow_name': 'train_model'}}
-
-    # Mock result without run_dir
-    mock_train_result.run_dir = None
-
-    workflow_mock.execute_activity_method = AsyncMock(
-        side_effect=[
-            None,  # delete_file_from_minio
-            None,  # update_experiment_run
-        ]
-    )
-
-    await train_model_workflow._cleanup_resources(
-        saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
-    )
-
-    # Verify cleanup_run_directory was NOT called (no run_dir)
-    # Only delete_file_from_minio + update_experiment_run
-    assert workflow_mock.execute_activity_method.call_count == 2
-
-
-# ============================================================================
-# Tests for _update_experiment_run()
-# ============================================================================
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_update_experiment_run_status_only(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test updating experiment run with status only."""
-    from model_manager.activities.experiment_tracking import UpdateType
-
-    metadata = {'metadata': {'experiment_run_id': 123}}
-
-    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
-
-    await train_model_workflow._update_experiment_run(
-        metadata=metadata,
-        experiment_run_id=123,
-        update_type=UpdateType.STATUS,
-        status=ExperimentStatus.TRAINING_SUCCESS,
-    )
-
-    workflow_mock.execute_activity_method.assert_called_once()
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_update_experiment_run_with_error(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test updating experiment run with error message."""
-    from model_manager.activities.experiment_tracking import UpdateType
-
-    metadata = {'metadata': {'experiment_run_id': 123}}
-
-    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
-
-    await train_model_workflow._update_experiment_run(
-        metadata=metadata,
-        experiment_run_id=123,
-        update_type=UpdateType.STATUS_WITH_ERROR,
-        status=ExperimentStatus.TRAINING_ERROR,
-        error_message='Training failed',
-    )
-
-    workflow_mock.execute_activity_method.assert_called_once()
-    call_args = workflow_mock.execute_activity_method.call_args[0][1]
-    assert call_args['error_message'] == 'Training failed'
-
-
-@mark.asyncio
-@patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-async def test_update_experiment_run_with_run_name(
-    workflow_mock: AsyncMock, train_model_workflow: TrainModel
-):
-    """Test updating experiment run with run_name."""
-    from model_manager.activities.experiment_tracking import UpdateType
-
-    metadata = {'metadata': {'experiment_run_id': 123}}
-
-    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
-
-    await train_model_workflow._update_experiment_run(
-        metadata=metadata,
-        experiment_run_id=123,
-        update_type=UpdateType.MODEL_SAVED,
-        status=ExperimentStatus.MLFLOW_SENT,
-        run_name='test_experiment-1',
-    )
-
-    workflow_mock.execute_activity_method.assert_called_once()
-    call_args = workflow_mock.execute_activity_method.call_args[0][1]
-    assert call_args['run_name'] == 'test_experiment-1'

todo-list.txt

@@ -0,0 +1,3 @@
+- remover os testes das classes alteradas e refazer de novo depois
+- alterar todos os comentários dos métodos que foram alterados
+- no final alterar o readme