Merge pull request #24 from Aignosi/SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious

SIENTIAPDE-1231, SIENTIAPDE-1222, SIENTIAPDE-1214: Enhance MLFlow, Model Repository, Logging, and Configuration

.env.example

@@ -27,3 +27,9 @@ MONGODB_PASSWORD="mongo_db_password"
 MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
 MONGODB_DATABASE="sientia"
 MONGODB_TTL_INDEX_HOURS="1"
+
+MINIO_ENDPOINT_URL="http://localhost:9000"
+MINIO_ACCESS_KEY="sientia"
+MINIO_SECRET_KEY="sientia"
+MINIO_REGION_NAME="sa-east-1"
+MINIO_DEFAULT_BUCKET="sientia"

.github/workflows/quality-gate.yml

@@ -30,12 +30,12 @@ jobs:
           owner: 'Aignosi'
           repositories: 'sientia-dataops-library,sientia-mlops-library'
 
-      - name: Prepare requirements.txt
+      - name: Prepare requirements-light.txt
         id: prepare-requirements
         run: |
           sed -e "s|git+ssh://git@github.com/|git+https://github.com/|g" \
               -e "s|git@github.com:|git+https://github.com/|g" \
-              requirements.txt > requirements_prepared.txt
+              requirements-light.txt > requirements_prepared.txt
           echo "PROCESSED_REQUIREMENTS_FILE=requirements_prepared.txt" >> $GITHUB_OUTPUT
 
       - name: Configure Git to use App Token

.gitignore

@@ -37,6 +37,7 @@ __pycache__/
 # Ignorar coverage
 htmlcov/
 .coverage
+coverage.xml
 
 # git keys
 git_key*
@@ -46,3 +47,8 @@ git_log
 .env
 
 tmp/
+catboost_info/
+
+.ruff_cache/
+.mypy_cache/
+mlruns/

README.md

@@ -1,42 +1,108 @@
 # Sientia DataOps Laborious
 
-A high-performance, scalable machine learning prediction system built on Temporal.io for industrial data processing and ML model inference. The Laborious system provides enterprise-grade ML model management, batch prediction processing, and real-time data export capabilities with comprehensive data quality validation and monitoring.
+A comprehensive, Temporal-based ML orchestration system for industrial data processing and model inference. Laborious delivers enterprise-grade batch prediction, model management, optional real-time export (OPC), and automated retraining with strong data quality validation and observability.
+
+## 📑 Table of Contents
+
+- [Features](#features)
+  - [Core Functionality](#core-functionality)
+  - [Advanced Capabilities](#advanced-capabilities)
+  - [Development & Quality Assurance](#development--quality-assurance)
+- [Architecture](#architecture)
+  - [Architecture Principles](#architecture-principles)
+  - [Key Components](#key-components)
+  - [Data Flow Architecture](#data-flow-architecture)
+  - [Security Architecture](#security-architecture)
+- [Workflows](#workflows)
+  - [Predictions Batch Workflow](#1-predictions-batch-workflow-predictions_batchpy)
+  - [Prediction Process Workflow](#2-prediction-process-workflow-prediction_processpy)
+  - [Format and Export Prediction Workflow](#3-format-and-export-prediction-workflow-format_and_export_predictionpy)
+  - [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy)
+- [Installation & Setup](#installation--setup)
+  - [Prerequisites](#prerequisites)
+  - [Environment Setup](#environment-setup)
+  - [Temporal Namespace Setup](#temporal-namespace-setup)
+  - [Local Development Setup](#local-development-setup)
+- [How to Run](#how-to-run)
+  - [Running the Laborious Application](#running-the-laborious-application)
+  - [Running Tests and Coverage](#running-tests-and-coverage)
+  - [Manual Test Execution](#manual-test-execution)
+  - [Manual Application Execution](#manual-application-execution)
+- [Code Quality & Validation](#code-quality--validation)
+  - [Overview](#overview)
+  - [Validation Tools](#validation-tools)
+  - [Tools Installation](#tools-installation)
+  - [Complete Validation](#complete-validation)
+  - [Automatic Fixes](#automatic-fixes)
+  - [Configuration](#configuration)
+  - [CI/CD Integration](#cicd-integration)
+  - [Best Practices](#best-practices)
+- [Testing](#testing)
+  - [Test Structure](#test-structure)
+  - [Test Execution](#test-execution)
+- [Monitoring and Metrics](#monitoring-and-metrics)
+  - [Application Health Metrics](#application-health-metrics)
+  - [Prediction Operation Metrics](#prediction-operation-metrics)
+  - [OPC Export Metrics](#opc-export-metrics)
+  - [Data Quality Metrics](#data-quality-metrics)
+- [Configuration](#configuration-1)
+  - [Environment Variables](#environment-variables)
+  - [OPC Configuration](#opc-configuration)
+  - [Workflow Configuration](#workflow-configuration)
+- [Development](#development)
+  - [Project Structure](#project-structure)
+  - [Adding New Features](#adding-new-features)
+- [Troubleshooting](#troubleshooting)
+  - [Common Issues](#common-issues)
+  - [Debug Mode](#debug-mode)
+- [Performance Tuning](#performance-tuning)
+  - [Key Parameters](#key-parameters)
+  - [Scaling Considerations](#scaling-considerations)
+- [Contributing](#contributing)
+  - [Code Quality Standards](#code-quality-standards)
+- [License](#license)
+- [Support](#support)
 
 ## Features
 
 ### Core Functionality
-- **Batch Prediction Processing**: High-throughput ML model inference using MLFlow models
+- **Batch Prediction Processing**: High-throughput ML inference using MLFlow models
-- **Temporal Workflow Orchestration**: Robust workflow management with automatic retry policies and fault tolerance
+- **Temporal Workflow Orchestration**: Robust workflow management with retries and fault tolerance
-- **Data Quality Gates**: Configurable filtering for data validation, MLFlow API responses, and custom validation rules
+- **Data Quality Gates**: Configurable filtering for input data and MLFlow API responses
-- **Multi-Model Support**: Flexible ML model management with retention policies and versioning
+- **Multi-Model Support**: Flexible model management with retention and versioning
-- **Real-time Data Export**: PostgreSQL persistence and OPC server integration for industrial systems
+- **Optional Real-time Export**: PostgreSQL persistence and OPC server integration for industrial systems
-- **Comprehensive Monitoring**: Prometheus metrics and detailed logging for operational visibility
+- **Comprehensive Monitoring**: Prometheus metrics and structured logging for observability
 
 ### Advanced Capabilities
-- **Incremental Data Processing**: Timestamp-based data loading to avoid reprocessing
+- **Incremental Data Processing**: Timestamp-based loading to avoid reprocessing
 - **Configurable Data Retention**: Model retention policies with automatic cleanup
-- **Notification System**: Integrated alerting and notification management via MongoDB
+- **Notification System**: Integrated alerting via MongoDB
-- **Scalable Architecture**: Kubernetes-ready deployment with horizontal scaling support
+- **Scalable Architecture**: Kubernetes-ready with horizontal scaling
-- **Model Retraining**: Automated model retraining workflows with production model updates
+- **Model Retraining**: Automated retraining workflows with production model updates
+
+### Development & Quality Assurance
+- **Code Quality Tools**: Ruff (lint/format), mypy (types), Bandit (security)
+- **Automated Validation**: `validate.sh` and CI quality gates
+- **Comprehensive Testing**: pytest with async support and high coverage
+- **Type Safety**: Static type checking with mypy
+- **Coverage Visualization**: Coverage Gutters integration
 
 ## Architecture
 
-The Laborious system uses a Temporal-based workflow architecture with clear separation of concerns and robust error handling. The architecture is designed for high availability, scalability, and operational excellence in production ML environments.
+Laborious uses a Temporal-based architecture with strong separation of concerns and defensive error handling for production ML.
-
 
 ### Architecture Principles
 
 #### 1. **Separation of Concerns**
-- **Worker Layer**: Manages Temporal workers, task queues, and application lifecycle
+- **Worker Layer**: Temporal workers, task queues, lifecycle
-- **Workflow Layer**: Orchestrates business logic and process coordination
+- **Workflow Layer**: Business orchestration and coordination
-- **Activity Layer**: Implements specific operations and external system interactions
+- **Activity Layer**: External system interactions and isolated operations
-- **Data Layer**: Handles data persistence, caching, and external service connections
+- **Data Layer**: Persistence, caching, connectors
 
 #### 2. **Fault Tolerance & Resilience**
-- **Automatic Retry Policies**: Configurable retry strategies for transient failures
+- **Automatic Retry Policies** for transient failures
-- **Circuit Breaker Pattern**: Prevents cascading failures in external service calls
+- **Graceful Degradation** and circuit breaking for dependencies
-- **Graceful Degradation**: System continues operating with reduced functionality
+- **Detailed Error Handling** with notifications
-- **Comprehensive Error Handling**: Detailed error reporting and notification integration
 
 #### 3. **Scalability & Performance**
 - **Horizontal Scaling**: Multiple worker instances for load distribution
@@ -53,66 +119,34 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa
 ### Key Components
 
 #### **Worker (`laborious/worker/worker.py`)**
-- **Purpose**: Main application orchestrator managing Temporal workers and task queues
+- Temporal client setup, worker lifecycle, task queues
-- **Responsibilities**:
+- Metrics server initialization, notification handler setup
-  - Temporal client initialization and connection management
+- Graceful shutdown and autoscaling-friendly behavior
-  - Worker lifecycle management and graceful shutdown
-  - Task queue configuration and load balancing
-  - Prometheus metrics server initialization
-  - Notification handler setup and configuration
-  - OPC server connection management
-- **Key Features**:
-  - Automatic scaling with `PollerBehaviorAutoscaling`
-  - Health check endpoints for Kubernetes liveness/readiness probes
-  - Graceful shutdown with cleanup procedures
-  - Multi-instance deployment support
-  - Two dedicated task queues: `predictions_batch-queue` and `minimal_retrain-queue`
 
 #### **Workflows (`laborious/workflows/`)**
-- **PredictionsBatch**: Main entry point for batch prediction pipelines
+- `predictions_batch.py`: Batch prediction entry point
-- **PredictionProcess**: Core prediction pipeline with MLFlow integration
+- `sub_workflows/prediction_process.py`: Core prediction pipeline
-- **FormatAndExportPrediction**: Data formatting and export operations
+- `sub_workflows/format_and_export_prediction.py`: Formatting and export
-- **MinimalRetrain**: Automated model retraining and deployment
+- `minimal_retrain.py`: Automated model retraining and production update
-- **Key Features**:
-  - Temporal workflow definitions with retry policies
-  - Child workflow orchestration and delegation
-  - Comprehensive error handling and recovery
-  - Configurable timeout and retry strategies
 
 #### **Activities (`laborious/activities/`)**
-- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance
+- `gates.py`: Data quality validation and filtering
-- **Gates**: Data quality validation and filtering mechanisms
+- `mlflow.py`: Transform and predict operations
-- **MLFlow**: Model transformation and prediction operations
+- `opc.py`: OPC UA export to industrial systems (optional)
-- **OPC**: Real-time data export to industrial OPC servers
+- `activities.py`: Aggregates activity interfaces
-- **Key Features**:
-  - Multiple inheritance pattern for unified activity interface
-  - Configurable filter policies and validation rules
-  - MLFlow model serving integration with configurable flavors
-  - OPC UA client with certificate-based authentication
-  - Comprehensive error handling and notification integration
-  - Support for multiple OPC servers with independent configurations
 
 #### **Data Services (`laborious/utils/`)**
-- **Connectors Config**: Environment variable-based configuration management
+- `connectors_config.py`: Env-driven configuration builders
-- **Repository**: Data access layer for MLFlow and OPC operations
+- `repository/model_repository.py`: MLFlow operations and retraining
-  - `model_repository.py`: MLFlow model operations and retraining
+- `repository/opc_repository.py`: OPC communication and writes
-  - `opc_repository.py`: OPC server communication and data writing
+- `filters/conditional_filters.py` and `filters/mlflow_filters.py`
-- **Filters**: Data quality validation and MLFlow response filtering
-  - `conditional_filters.py`: Input data validation filters
-  - `mlflow_filters.py`: MLFlow API response validation filters
-- **Key Features**:
-  - Environment variable-based configuration with sensible defaults
-  - Connection pool management and optimization
-  - Security credential management
-  - Configuration validation and error handling
-  - Support for multiple OPC servers and MLFlow model flavors
 
 ### Data Flow Architecture
 
 #### **1. Batch Prediction Pipeline**
 ```
 Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
-MLFlow Prediction → Response Validation → Export (PostgreSQL + OPC)
+MLFlow Prediction → Response Validation → Export (PostgreSQL [+ OPC])
 ```
 
 #### **2. Model Retraining Pipeline**
@@ -121,33 +155,21 @@ Training Data → Model Retraining → Quality Validation →
 Production Update → Notification & Monitoring
 ```
 
-#### **3. Real-time Export Pipeline**
-```
-Prediction Results → Data Formatting → OPC Server Write → 
-Success/Failure Metrics → Notification System
-```
-
 ### Security Architecture
 
 #### **Authentication & Authorization**
-- **Certificate-based OPC Authentication**: Secure industrial communication
+- **MLFlow API Authentication**: Username/password
-- **MLFlow API Authentication**: Username/password with secure transmission
+- **Database Security**: Encrypted connections and credential management
-- **Database Connection Security**: Encrypted connections with credential management
+- **OPC Certificates** (if enabled): Client/server certs
-- **Kubernetes Secrets Integration**: Secure credential storage and access
+- **Kubernetes Secrets**: Secure secret storage
 
 #### **Network Security**
-- **TLS/SSL Encryption**: Secure communication channels
+- TLS/SSL, network policies, service mesh, firewalls, VPN
-- **Network Isolation**: Kubernetes network policies and service mesh
-- **Firewall Rules**: Controlled access to external services
-- **VPN Integration**: Secure remote access and management
 
 #### **Data Security**
-- **Data Encryption**: At-rest and in-transit encryption
+- At-rest/in-transit encryption, RBAC, audit logging, lifecycle management
-- **Access Control**: Role-based access control (RBAC)
-- **Audit Logging**: Comprehensive access and operation logging
-- **Data Retention**: Configurable data lifecycle management
 
-## 🔄 Workflows
+## Workflows
 
 ### 1. Predictions Batch Workflow (`predictions_batch.py`)
 
@@ -351,8 +373,9 @@ flowchart LR
 - Temporal server/cluster
 - PostgreSQL database
 - MLFlow server
-- OPC server(s)
+- MinIO object storage (for MLFlow artifacts)
 - MongoDB server (for notifications)
+- OPC server(s) if using OPC export
 
 **Note**: External dependencies must be available either through:
 - Kubernetes cluster deployment
@@ -488,6 +511,68 @@ fi
 python -m laborious.worker.worker
 ```
 
+## Code Quality & Validation
+
+### Overview
+
+Since Python is not compiled, we validate quality, security, and correctness before execution.
+
+### Validation Tools
+
+- Ruff: Linting and formatting
+- mypy: Static type checking
+- Bandit: Security analysis
+- pytest: Unit/integration testing with coverage
+
+### Tools Installation
+
+```bash
+pip install -r requirements-dev.txt
+```
+
+### Complete Validation
+
+Option 1 (recommended):
+```bash
+./validate.sh
+```
+The script runs, in order:
+1. Format check (Ruff)
+2. Linting (Ruff)
+3. Type checking (mypy)
+4. Security analysis (Bandit)
+5. Tests with coverage (pytest)
+
+Option 2 (individual commands):
+```bash
+ruff format --check laborious/ tests/
+ruff check laborious/ tests/
+mypy laborious/
+bandit -r laborious/ -ll
+pytest tests/ --cov=laborious --cov-report=term-missing
+```
+
+### Automatic Fixes
+
+```bash
+ruff format laborious/ tests/
+ruff check --fix laborious/ tests/
+```
+
+### Configuration
+
+All settings reside in `pyproject.toml` (Ruff, mypy, pytest, Bandit).
+
+### CI/CD Integration
+
+The workflow at `.github/workflows/quality-gate.yml` executes validations on each push/PR.
+
+### Best Practices
+
+- Run `./validate.sh` before committing
+- Use `ruff check --watch` for continuous feedback
+- Add type hints and tests for new code
+
 ## 🧪 Testing
 
 ### Test Structure

encode.sh

@@ -0,0 +1,13 @@
+source ./venv/bin/activate
+
+pip install pathspec
+pip install pyyaml
+
+echo "
+.git" >> .gitignore
+
+python encrypt.py ./ code --ignore .gitignore --chunk-size 100000
+
+sed -i '/.git/d' .gitignore
+
+xdg-open .

encrypt.py

@@ -0,0 +1,112 @@
+import os
+import argparse
+from pathspec import PathSpec
+import yaml
+
+'''
+Usage:
+    python .\encrypt.py path_to_dir output_file --ignore ignore_file --chunk-size 100000
+'''
+
+
+def load_ignore_patterns(ignore_file, include_library):
+    # Ensure the .gitignore file exists
+    if not os.path.exists(ignore_file):
+        raise FileNotFoundError(f"Ignore file not found at {ignore_file}")
+
+    # Load and parse the .gitignore patterns
+    with open(ignore_file, 'r') as file:
+        patterns = file.readlines()
+        if not include_library:
+            patterns.append('**/deploy/library/')
+
+        spec = PathSpec.from_lines('gitwildmatch', patterns)
+        return spec
+
+
+def is_ignored(file_path, spec):
+    """Check if a file should be ignored based on the ignore patterns."""
+    return spec.match_file(file_path) if spec else False
+
+
+def encode_file_tree_to_yaml(directory, ignore_file, include_library):
+    """Encode the file tree into a single YAML file."""
+    ignore_patterns = load_ignore_patterns(
+        ignore_file, include_library) if ignore_file else None
+    file_tree = {}
+
+    for root, dirs, files in os.walk(directory):
+        # Skip ignored directories
+        dirs[:] = [d for d in dirs if not is_ignored(
+            os.path.join(root, d), ignore_patterns)]
+
+        for file in files:
+            file_path = os.path.join(root, file)
+
+            # Skip ignored files
+            if is_ignored(file_path, ignore_patterns):
+                continue
+
+            # Read file content
+            try:
+                with open(file_path, 'r', encoding='utf-8') as f:
+                    content = f.read()
+            except Exception as e:
+                print(f"Error reading file {file_path}: {e}")
+                raise
+
+            # Create nested dictionary structure
+            path_parts = os.path.relpath(file_path, directory).split(os.sep)
+            current_level = file_tree
+
+            # all except the last part (the file name)
+            for part in path_parts[:-1]:
+                current_level = current_level.setdefault(part, {})
+
+            # Add the file and its content
+            current_level[path_parts[-1]] = content
+    return yaml.dump(file_tree, default_flow_style=False)
+
+
+def chunk_and_write_file_tree_to_yaml(yaml_content, output_file, chunk_size=None):
+    """Chunk the YAML content and write it to the output file."""
+
+    chunks = [yaml_content] if chunk_size is None else [
+        yaml_content[i:i + chunk_size] for i in range(0, len(yaml_content), chunk_size)]
+
+    for i, chunk in enumerate(chunks):
+        chunk_file = f"{output_file}_{i}.yaml"
+        # Write the file tree to the output YAML file
+        with open(chunk_file, 'w', encoding='utf-8') as yaml_file:
+            yaml_file.write(chunk)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Encrypts file tree to yaml file")
+    parser.add_argument("input_directory", help="Directory to encode")
+    parser.add_argument("output_yaml_file", help="Output YAML file")
+    parser.add_argument("--ignore", default=None,
+                        help="Path to the ignore file")
+    parser.add_argument("--chunk-size", type=int, default=None,
+                        help="Chunk size for the output YAML file")
+    parser.add_argument("--library", type=bool, default=False,
+                        help="Incude the library in the output YAML file")
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Example usage
+    directory_to_encode = args.input_directory
+    ignore_file_path = args.ignore
+    output_yaml_file = args.output_yaml_file
+    include_library = args.library
+
+    content = encode_file_tree_to_yaml(
+        directory_to_encode, ignore_file_path, include_library)
+    chunk_and_write_file_tree_to_yaml(
+        content, output_yaml_file, args.chunk_size)
+
+
+if __name__ == "__main__":
+    main()

laborious/activities/activities.py

@@ -1,16 +1,18 @@
-from temporalio import activity, workflow
+from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from sientia_do.temporal.activities.postgres import Postgres
-    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
-    from sientia_do.observability.logger import Logger
-    from laborious.activities.mlflow import MLFlow
-    from laborious.activities.gates import Gates
-    from laborious.activities.opc import OPC
     from typing import Any
 
+    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+    from sientia_do.observability.logger import Logger
 
-class Activities(Postgres, MLFlow, Gates, OPC):
+    from laborious.activities.gates import Gates
+    from laborious.activities.mlflow import MLFlow
+    from laborious.activities.opc import OPC
+    from laborious.activities.storage import Storage
+
+
+class Activities(Storage, MLFlow, Gates, OPC):
     """
     Main activities orchestrator for the Laborious system.
 
@@ -32,12 +34,15 @@ class Activities(Postgres, MLFlow, Gates, OPC):
         notification_handler (NotificationHandler): Notification management instance
     """
 
-    def __init__(self,
+    def __init__(
+        self,
         postgres_config: dict[str, Any],
         mlflow_config: dict[str, Any],
+        minio_config: dict[str, Any],
         opc_config: dict[str, Any],
         logger: Logger,
-                 notification_handler: NotificationHandler):
+        notification_handler: NotificationHandler,
+    ):
         """
         Initialize the Activities orchestrator with all required configurations.
 
@@ -58,30 +63,36 @@ class Activities(Postgres, MLFlow, Gates, OPC):
             Exception: If any parent class initialization fails
         """
         # Initialize parent classes
-        Postgres.__init__(self, host=postgres_config['host'],
+        Storage.__init__(
+            self,
+            host=postgres_config['host'],
             port=postgres_config['port'],
             user=postgres_config['user'],
             password=postgres_config['password'],
             dbname=postgres_config['dbname'],
             min_connections=postgres_config['min_connections'],
             max_connections=postgres_config['max_connections'],
+            minio_config=minio_config,
             logger=logger,
-                          notification_handler=notification_handler)
+            notification_handler=notification_handler,
+        )
 
-        MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
+        MLFlow.__init__(
+            self,
+            mlflow_host=mlflow_config['host'],
             mlflow_port=mlflow_config['port'],
             mlflow_username=mlflow_config['username'],
             mlflow_password=mlflow_config['password'],
+            minio_config=minio_config,
             logger=logger,
-                        notification_handler=notification_handler)
+            notification_handler=notification_handler,
+        )
 
-        Gates.__init__(self, logger=logger,
+        Gates.__init__(self, logger=logger, notification_handler=notification_handler)
-                       notification_handler=notification_handler)
 
-        OPC.__init__(self,
+        OPC.__init__(
-                     opc_servers=opc_config,
+            self, opc_servers=opc_config, logger=logger, notification_handler=notification_handler
-                     logger=logger,
+        )
-                     notification_handler=notification_handler)
 
     async def shutdown(self):
         """
@@ -95,5 +106,5 @@ class Activities(Postgres, MLFlow, Gates, OPC):
         The method should be called before the application terminates to ensure
         proper resource cleanup and prevent resource leaks.
         """
-        Postgres.close(self)
+        Storage.close(self)
         await OPC.shutdown(self)

laborious/activities/gates.py

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

laborious/activities/mlflow.py

@@ -1,20 +1,25 @@
 from temporalio import activity, workflow
 
-
 with workflow.unsafe.imports_passed_through():
-    from datetime import datetime
+    import traceback
-    from pandas import Timestamp, to_datetime
+    from typing import Any
-    from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
-    from sientia_do.temporal.activities.base import BaseActivity
+    import numpy as np
+    from pandas import DataFrame, to_datetime
+    from sientia_do.formatters import create_sample_dict
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
     from sientia_do.observability.logger import Logger
-    from sientia_do.formatters import create_sample_dict
+    from sientia_do.temporal.activities.base import BaseActivity
+    from sientia_do.temporal.constants import (
+        DATETIME_FORMAT,
+        DATETIME_FORMAT_MS_WITH_TZ,
+        DATETIME_FORMAT_WITH_TZ,
+        now,
+    )
+
+    from laborious.utils.repository.minio_repository import MinioRepository
     from laborious.utils.repository.model_repository import MLFlowRepository
-    from typing import Any
-    import numpy as np
-    from pandas import DataFrame
-    import traceback
 
 
 class MLFlow(BaseActivity):
@@ -36,8 +41,16 @@ class MLFlow(BaseActivity):
         model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
     """
 
-    def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
+    def __init__(
-                 mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
+        self,
+        mlflow_host: str,
+        mlflow_port: int,
+        mlflow_username: str,
+        minio_config: dict[str, Any],
+        mlflow_password: str,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
         """
         Initialize MLFlow activities with server configuration.
 
@@ -52,18 +65,31 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If MLFlowRepository initialization fails
         """
-        BaseActivity.__init__(
+        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
-            self, logger, notification_handler, set_error_counter=True)
         self.mlflow_host = mlflow_host
         self.mlflow_port = mlflow_port
         self.mlflow_username = mlflow_username
         self.mlflow_password = mlflow_password
 
         self.model_monitoring_repository = MLFlowRepository(
-            f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger
+            f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
         )
 
-    @activity.defn(name="request_transform")
+        if not hasattr(self, 'minio_repository'):
+            self.minio_repository: MinioRepository | None = None
+
+        if self.minio_repository is None:
+            self.minio_repository = MinioRepository(
+                logger=logger,
+                notification_handler=notification_handler,
+                minio_endpoint_url=minio_config['endpoint_url'],
+                minio_access_key=minio_config['access_key'],
+                minio_secret_key=minio_config['secret_key'],
+                minio_region_name=minio_config['region_name'],
+                minio_default_bucket=minio_config['default_bucket'],
+            )
+
+    @activity.defn(name='request_transform')
     async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Transform input data using MLFlow models.
@@ -99,7 +125,7 @@ class MLFlow(BaseActivity):
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug("Raw input data:", metadata)
+        self.debug('Raw input data:', metadata)
         self.debug(data.head(5).to_string(), metadata)
 
         # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
@@ -108,14 +134,12 @@ class MLFlow(BaseActivity):
         )
 
         # Pivot data for model input format
-        data = data.pivot(
+        data = data.pivot(index='timestamp', columns='variable', values='value')
-            index='timestamp', columns='variable',
-            values='value')
         data.fillna(np.nan, inplace=True)
         # data.reset_index(inplace=True)
         data.columns.name = None
 
-        self.debug("Processed input data:", metadata)
+        self.debug('Processed input data:', metadata)
         self.debug(data.head(5).to_string(), metadata)
 
         # Request transformation from MLFlow model
@@ -124,16 +148,20 @@ class MLFlow(BaseActivity):
         )
 
         self.debug(
-            f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
         self.debug(
-            f"Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data transformed successfully", metadata)
+        self.info('Data transformed successfully', metadata)
 
         return response_data
 
-    @activity.defn(name="request_predict")
+    @activity.defn(name='request_predict')
     async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Execute predictions using MLFlow models.
@@ -169,14 +197,15 @@ class MLFlow(BaseActivity):
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug(f"Input data for: \n {data.head(5).to_string()}", metadata)
+        self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
 
         # Convert numpy.nan to None for model compatibility
         data.replace(np.nan, None, inplace=True)
 
         data['timestamp'] = data.index
         data['timestamp'] = to_datetime(
-            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ).dt.strftime(DATETIME_FORMAT)
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
+        ).dt.strftime(DATETIME_FORMAT)
 
         # Request prediction from MLFlow model
         response_data = self.model_monitoring_repository.predict(
@@ -184,13 +213,15 @@ class MLFlow(BaseActivity):
         )
 
         self.debug(
-            f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
+            f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
+            metadata,
+        )
 
-        self.info("Data predicted successfully", metadata)
+        self.info('Data predicted successfully', metadata)
 
         return response_data
 
-    @activity.defn(name="retrain_model")
+    @activity.defn(name='retrain_model')
     async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Retrain MLFlow models with updated training data.
@@ -222,51 +253,88 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If retraining fails or encounters critical errors
         """
+
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
+
         metadata = input_data['metadata']
-        data = DataFrame(input_data['data'])
+        object_key = input_data['object_key']
+
+        self.info(f'Loading retrain data from Key: {object_key}', metadata)
+
+        try:
+            data = self.minio_repository.get_parquet_as_dataframe(
+                object_key=object_key, metadata=metadata
+            )
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.send_notification(
+                metadata=metadata,
+                notification_id='ERROR_LOADING_RETRAIN_DATA',
+                message=f'Error loading retrain data: {e}',
+                block='retrain_model',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+            self.error(trace, metadata)
+            return {
+                'success': False,
+                'message': f'Error loading retrain data: {e}',
+                'traceback': trace,
+                'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
+            }
+
+        self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
+
         model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
 
         self.info(f'Retraining model {model_name}...', metadata)
 
         timestamp = data['timestamp'].max()
         self.debug(f'Timestamp: {timestamp}', metadata)
 
+        # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
+        data = data.sort_values('created_at', ascending=False).drop_duplicates(
+            subset=['variable', 'timestamp'], keep='first'
+        )
+
         data.drop(columns=['model_id'], inplace=True, errors='ignore')
         data.drop(columns=['created_at'], inplace=True, errors='ignore')
 
-        data = data.pivot(index='timestamp', columns='variable',
+        # Pivot data for model input format
-                          values='value')
+        data = data.pivot(index='timestamp', columns='variable', values='value')
-        data.sort_index(inplace=True)
+        data.fillna(np.nan, inplace=True)
-        data.reset_index(inplace=True)
+        # data.reset_index(inplace=True)
-
-        data = data.dropna()
         data.columns.name = None
 
-        try:
+        data['timestamp'] = data.index
-            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
+        data['timestamp'] = to_datetime(
-                data=data,
+            data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
-                model_name=model_name
+        ).dt.strftime(DATETIME_FORMAT)
+        data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
+
+        data.columns.name = None
+
+        retrain_output = self.model_monitoring_repository.retrain_model(
+            data=data, model_name=model_name, model_config=model_config, metadata=metadata
         )
 
-            return {
+        if not retrain_output['success']:
-                'status': retrain_output,
+            trace = retrain_output['traceback']
-                'timestamp': timestamp,
-                'experiment': experiment
-            }
-        except Exception as e:
-            trace = traceback.format_exc()
             self.send_notification(
                 metadata=metadata,
                 notification_id='RETRAIN_MODEL_ERROR',
-                message=f'Error retraining model {model_name}: {e}',
+                message=f'Error retraining model {model_name}: {retrain_output["message"]}',
                 block='retrain_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             self.error(trace, metadata=metadata)
-            raise e
 
-    @activity.defn(name="update_production_model")
+        return {**retrain_output, 'timestamp': timestamp}
+
+    @activity.defn(name='update_production_model')
     async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
         """
         Update production model with newly trained model version.
@@ -305,29 +373,18 @@ class MLFlow(BaseActivity):
         """
         metadata = input_data['metadata']
         model_name = input_data['model_name']
-        model_id = input_data['model_id']
         experiment = input_data['experiment']
-        timestamp = input_data['timestamp']
-        status = input_data['status']
-
         self.info(
-            f'Updating production model {model_name} from experiment {experiment}...', metadata)
+            f'Updating production model {model_name} from experiment {experiment}...', metadata
+        )
 
         try:
             response = self.model_monitoring_repository.update_production_model(
-                experiment=experiment,
+                experiment=experiment, model_name=model_name, metadata=metadata
-                model_name=model_name
             )
 
-            report = DataFrame([response])
+            self.info(f'Production model {model_name} updated successfully', metadata)
-            report['model_id'] = model_id
+            return response
-            report['model_name'] = model_name
-            report['timestamp'] = timestamp
-            report['status'] = status
-
-            self.info(
-                f'Production model {model_name} updated successfully', metadata)
-            return report.to_dict()
 
         except Exception as e:
             trace = traceback.format_exc()
@@ -337,7 +394,7 @@ class MLFlow(BaseActivity):
                 message=f'Error updating production model {model_name}: {e}',
                 block='update_production_model',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             self.error(trace, metadata=metadata)
             raise e

laborious/activities/opc.py

@@ -1,15 +1,16 @@
 from temporalio import activity, workflow
 
-
 with workflow.unsafe.imports_passed_through():
+    import traceback
+    from typing import Any
+
+    from pandas import DataFrame
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
-    from sientia_do.temporal.activities.base import BaseActivity
     from sientia_do.observability.logger import Logger
+    from sientia_do.temporal.activities.base import BaseActivity
+
     from laborious.utils.repository.opc_repository import OpcRepository
-    from typing import Any
-    import traceback
-    from pandas import DataFrame
 
 OPC_WRITTING_ERROR_CONFIDENCE = 12
 
@@ -33,15 +34,17 @@ class OPC(BaseActivity):
         notification_handler (NotificationHandler): Notification management instance
     """
 
-    def __init__(self, opc_servers: dict[str, dict[str, Any]],
+    def __init__(
-                 logger: Logger, notification_handler: NotificationHandler):
+        self,
-
+        opc_servers: dict[str, dict[str, Any]],
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
         self.logger = logger
         self.notification_handler = notification_handler
         self.opc_servers = opc_servers
 
-        BaseActivity.__init__(
+        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
-            self, logger, notification_handler, set_error_counter=True)
 
         self.opc_repository: dict[str, OpcRepository] = {}
         self.opc_servers = opc_servers
@@ -70,10 +73,10 @@ class OPC(BaseActivity):
             the initialization of other OPC servers. Each server is handled
             independently to ensure maximum availability.
         """
-        self.logger.info("Initializing OPC servers...")
+        self.logger.info('Initializing OPC servers...')
-        for id, server in self.opc_servers.items():
+        for opc_id, server in self.opc_servers.items():
-            self.opc_repository[id] = OpcRepository(
+            self.opc_repository[opc_id] = OpcRepository(
-                id=server['id'],
+                opc_id=server['id'],
                 url=server['url'],
                 logger=self.logger,
                 server_uri=server['server_uri'],
@@ -82,30 +85,35 @@ class OPC(BaseActivity):
                 server_cert_path=server['server_cert_path'],
                 notification_handler=self.notification_handler,
                 reconnection_interval=server['reconnection_interval'],
-                pod_id=self.pod_id
+                pod_id=self.pod_id,
             )
-            is_connected, error_data = await self.opc_repository[id].connect()
+            is_connected, error_data = await self.opc_repository[opc_id].connect()
             if not is_connected:
                 self.send_notification(
                     metadata={
                         'model_id': '-',
                         'model_name': '-',
                         'workflow_name': '-',
-                        'schedule_name': 'INITIALIZATION'
+                        'schedule_name': 'INITIALIZATION',
                     },
                     notification_id=error_data['notification_id'],
                     message=error_data['message'],
                     block=error_data['block'],
                     level=error_data.get('level', NotificationLevel.ERROR),
-                    attachment_content=error_data.get(
+                    attachment_content=error_data.get('attachment_content', None),
-                        'attachment_content', None)
                 )
             else:
-                self.logger.info(
+                self.logger.info(f'OPC server {opc_id} connected successfully.')
-                    f"OPC server {id} connected successfully.")
 
-    async def write_data(self, server_id: str, tag: str, data: Any,
+    async def write_data(
-                         data_type: str, tag_type: str, metadata: dict[str, Any]) -> bool:
+        self,
+        server_id: str,
+        tag: str,
+        data: Any,
+        data_type: str,
+        tag_type: str,
+        metadata: dict[str, Any],
+    ) -> bool:
         """
         Write data to a specific OPC server tag with comprehensive error handling.
 
@@ -127,7 +135,8 @@ class OPC(BaseActivity):
 
         try:
             is_success, error_data = await self.opc_repository[server_id].write_data(
-                tag, data, data_type, self.logger, metadata)
+                tag, data, data_type, self.logger, metadata
+            )
             if not is_success:
                 self.send_notification(
                     metadata=metadata,
@@ -135,8 +144,7 @@ class OPC(BaseActivity):
                     message=error_data['message'],
                     block=error_data['block'],
                     level=error_data.get('level', NotificationLevel.ERROR),
-                    attachment_content=error_data.get(
+                    attachment_content=error_data.get('attachment_content', None),
-                        'attachment_content', None)
                 )
                 return False
             return True
@@ -144,11 +152,11 @@ class OPC(BaseActivity):
             trace = traceback.format_exc()
             self.send_notification(
                 metadata=metadata,
-                notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
+                notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
-                message=f"Error writing data to OPC server: {e}",
+                message=f'Error writing data to OPC server: {e}',
-                block="write_opc_data",
+                block='write_opc_data',
                 level=NotificationLevel.ERROR,
-                attachment_content=trace
+                attachment_content=trace,
             )
             raise e
 
@@ -174,21 +182,26 @@ class OPC(BaseActivity):
             This helps operators quickly identify configuration issues.
         """
         if self.opc_repository.get(server_id) is None:
-            message = f"OPC server {server_id} not found to perform write operation."
+            message = f'OPC server {server_id} not found to perform write operation.'
             self.send_notification(
                 metadata=metadata,
-                notification_id="OPC_SERVER_NOT_FOUND",
+                notification_id='OPC_SERVER_NOT_FOUND',
                 message=message,
-                block="write_opc_data",
+                block='write_opc_data',
                 level=NotificationLevel.ERROR,
-                attachment_content=f"OPC servers: {list(self.opc_repository.keys())}"
+                attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
             )
             return False
         return True
 
     async def manage_output_tags(
-            self, server_id: str, config: dict[str, Any], data: DataFrame,
+        self,
-            metadata: dict[str, Any], success: bool) -> tuple[bool, int]:
+        server_id: str,
+        config: dict[str, Any],
+        data: DataFrame,
+        metadata: dict[str, Any],
+        success: bool,
+    ) -> tuple[bool, int]:
         """
         Manage the writing of prediction and confidence data to OPC server tags.
 
@@ -225,11 +238,13 @@ class OPC(BaseActivity):
                     data=data.head(1)['prediction'].values[0],
                     data_type=tag_config['data_type'],
                     tag_type='prediction',
-                    metadata=metadata
+                    metadata=metadata,
                 )
                 if local_success:
                     self.info(
-                        f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata)
+                        f'Prediction data written to OPC server {server_id} for tag {tag}.',
+                        metadata,
+                    )
                     count += 1
                 success = success and local_success
 
@@ -241,11 +256,13 @@ class OPC(BaseActivity):
                     data=data.head(1)['prediction_confidence'].values[0],
                     data_type=tag_config['data_type'],
                     tag_type='confidence',
-                    metadata=metadata
+                    metadata=metadata,
                 )
                 if local_success:
                     self.info(
-                        f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata)
+                        f'Confidence data written to OPC server {server_id} for tag {tag}.',
+                        metadata,
+                    )
                     count += 1
                 success = success and local_success
 
@@ -271,29 +288,33 @@ class OPC(BaseActivity):
 
         """
         metadata = input_data['metadata']
-        self.info("Writing data to OPC servers...", metadata)
+        self.info('Writing data to OPC servers...', metadata)
         data = DataFrame(input_data['data'])
         opc_output_config = input_data['opc_output_config']
-        self.info(f"Data to write: {data.size} rows", metadata)
+        self.info(f'Data to write: {data.size} rows', metadata)
 
         success = True
 
         for server_id, config in opc_output_config.items():
-
             if not self.validate_server(server_id, metadata):
                 success = False
                 continue
 
             local_success, local_count = await self.manage_output_tags(
-                server_id, config, data, metadata, success)
+                server_id, config, data, metadata, success
+            )
             success = success and local_success
 
             self.info(
-                f"Process completed for OPC server {server_id}: {local_count} of {len(config['prediction_tags'])} prediction tags and {len(config['confidence_tags'])} confidence tags", metadata)
+                f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
+                metadata,
+            )
 
         return self.process_confidence(data, success, metadata)
 
-    def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]:
+    def process_confidence(
+        self, data: DataFrame, success: bool, metadata: dict[str, Any]
+    ) -> dict[Any, Any]:
         """
         Process prediction confidence based on OPC write operation success.
 
@@ -323,12 +344,12 @@ class OPC(BaseActivity):
         if not success:
             data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
             self.debug(
-                f"Some data could not be written to OPC servers, setting confidence to  {OPC_WRITTING_ERROR_CONFIDENCE}.",
+                f'Some data could not be written to OPC servers, setting confidence to  {OPC_WRITTING_ERROR_CONFIDENCE}.',
-                metadata
+                metadata,
             )
 
         else:
-            self.debug("Data written to OPC servers successfully.", metadata)
+            self.debug('Data written to OPC servers successfully.', metadata)
 
         return data.to_dict()
 

laborious/activities/storage.py

@@ -0,0 +1,137 @@
+from temporalio import activity, workflow
+
+with workflow.unsafe.imports_passed_through():
+    # Extend the Temporal Postgres activities for convenient query -> MinIO export
+    import traceback
+    from typing import Any
+
+    import pandas as pd
+    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.postgres import Postgres
+    from sientia_do.temporal.constants import now
+
+    from laborious.utils.repository.minio_repository import MinioRepository
+
+
+DATETIME_FILENAME_FORMAT = '%Y-%m-%d_%H-%M-%S'
+
+
+class Storage(Postgres):
+    """
+    Extensions for Postgres activities with a helper to export query results
+    directly to MinIO as Parquet and return the object name.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        user: str,
+        password: str,
+        dbname: str,
+        min_connections: int,
+        max_connections: int,
+        minio_config: dict[str, Any],
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
+        super().__init__(
+            host=host,
+            port=port,
+            user=user,
+            password=password,
+            dbname=dbname,
+            min_connections=min_connections,
+            max_connections=max_connections,
+            logger=logger,
+            notification_handler=notification_handler,
+        )
+
+        if not hasattr(self, 'minio_repository'):
+            self.minio_repository: MinioRepository | None = None
+
+        if self.minio_repository is None:
+            self.minio_repository = MinioRepository(
+                logger=logger,
+                notification_handler=notification_handler,
+                minio_endpoint_url=minio_config['endpoint_url'],
+                minio_access_key=minio_config['access_key'],
+                minio_secret_key=minio_config['secret_key'],
+                minio_region_name=minio_config['region_name'],
+                minio_default_bucket=minio_config['default_bucket'],
+            )
+
+    @activity.defn(name='query_to_minio')
+    async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Execute SQL query, write result as Parquet to MinIO, and return object name.
+
+        Args (input_data):
+            metadata (dict): Workflow metadata
+            query (str): SQL query
+            model_name (str): Model name for object naming
+            object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain)
+
+        Returns:
+            dict: { success: bool, object_name: str, uri: str }
+        """
+
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
+
+        metadata = input_data.get('metadata', {})
+        object_prefix = input_data.get('object_prefix', 'datasets/retrain')
+
+        timestamp = now().strftime(DATETIME_FILENAME_FORMAT)
+        object_name = f'{object_prefix}_{timestamp}.parquet'
+        uri = f's3://{self.minio_repository.minio_bucket}/{object_name}'
+
+        try:
+            data = await self.load_custom_query(input_data)
+            if not data:
+                self.error('query_to_minio failed: No data returned from query', metadata)
+                return {'success': False, 'message': 'No data returned from query'}
+
+            # Ensure we have a DataFrame
+            data = pd.DataFrame(data)
+
+            # Write parquet to memory and upload via persistent client
+            self.minio_repository.store_dataframe_as_parquet(
+                dataframe=data, uri=uri, object_name=object_name, metadata=metadata
+            )
+
+            return {'success': True, 'object_key': object_name, 'uri': uri}
+        except Exception as e:
+            trace = traceback.format_exc()
+            self.send_notification(
+                metadata=metadata,
+                notification_id='ERROR_STORING_QUERY_TO_MINIO',
+                message=f'Error storing query to MinIO: {e}',
+                block='query_to_minio',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+
+            self.error(trace, metadata)
+
+            return {'success': False, 'message': str(e)}
+
+    def close(self) -> None:
+        """Close Storage resources (MinIO client and Postgres engine)."""
+        try:
+            if hasattr(self, 'minio_repository') and self.minio_repository is not None:
+                try:
+                    self.minio_repository.close()
+                finally:
+                    self.minio_repository = None
+        finally:
+            # Ensure Postgres resources are disposed as well
+            try:
+                super().close()
+            except Exception:
+                self.logger.error('Error closing Postgres resources')
+
+    def __del__(self):
+        self.close()

laborious/metrics.py

@@ -23,50 +23,50 @@ Metric Labels:
 - opc_server_id: Identifier for OPC server operations
 """
 
-from prometheus_client import Gauge, Counter, Histogram
+from prometheus_client import Counter, Gauge, Histogram
 
 # Application health metric
 APP_UP = Gauge(
-    "app_up",
+    'app_up',
-    "Indicates if the application is running (1) or shutting down (0)",
+    'Indicates if the application is running (1) or shutting down (0)',
-    ["pod_id"],
+    ['pod_id'],
 )
 
 # Core labels used across multiple metrics
-CORE_LABELS = ["pod_id", "model_name", "pipeline_name"]
+CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
 
 # Prediction operation metrics
 PREDICTIONS_WRITTEN_COUNT = Counter(
-    "laborious_predictions_written_count",
+    'laborious_predictions_written_count',
-    "Number of predictions written to the database table predictions",
+    'Number of predictions written to the database table predictions',
     CORE_LABELS,
 )
 
 # Prediction quality metrics
 PREDICTION_CONFIDENCE_MONITOR = Gauge(
-    "laborious_prediction_confidence_monitor",
+    'laborious_prediction_confidence_monitor',
-    "Current confidence of each prediction",
+    'Current confidence of each prediction',
     CORE_LABELS,
 )
 
 # Performance monitoring metrics
 PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
-    "laborious_prediction_response_time_monitor",
+    'laborious_prediction_response_time_monitor',
-    "Current response time of each prediction",
+    'Current response time of each prediction',
     CORE_LABELS,
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )
 
 # OPC export metrics
 PREDICTION_OPC_WRITING_COUNT = Counter(
-    "laborious_prediction_opc_writing_count",
+    'laborious_prediction_opc_writing_count',
-    "Number of predictions written to the OPC server",
+    'Number of predictions written to the OPC server',
-    [*CORE_LABELS, "opc_server_id"],
+    [*CORE_LABELS, 'opc_server_id'],
 )
 
 PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
-    "laborious_prediction_opc_writing_response_time_monitor",
+    'laborious_prediction_opc_writing_response_time_monitor',
-    "Current response time of each prediction written to the OPC server",
+    'Current response time of each prediction written to the OPC server',
-    [*CORE_LABELS, "opc_server_id"],
+    [*CORE_LABELS, 'opc_server_id'],
-    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0]
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )

laborious/utils/connectors_config.py

@@ -1,9 +1,9 @@
-from os import getenv
 import json
-from typing import Dict, Any
+from os import getenv
+from typing import Any
 
 
-def build_postgres_config() -> Dict[str, Any]:
+def build_postgres_config() -> dict[str, Any]:
     """
     Build PostgreSQL database configuration from environment variables.
 
@@ -30,11 +30,11 @@ def build_postgres_config() -> Dict[str, Any]:
         'password': getenv('POSTGRES_PASSWORD', 'sientia'),
         'dbname': getenv('POSTGRES_DBNAME', 'sientia'),
         'min_connections': int(getenv('POSTGRES_MIN_CONNECTIONS', '5')),
-        'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20'))
+        'max_connections': int(getenv('POSTGRES_MAX_CONNECTIONS', '20')),
     }
 
 
-def build_mlflow_config() -> Dict[str, Any]:
+def build_mlflow_config() -> dict[str, Any]:
     """
     Build MLFlow server configuration from environment variables.
 
@@ -55,11 +55,11 @@ def build_mlflow_config() -> Dict[str, Any]:
         'host': getenv('MLFLOW_HOST', 'http://localhost'),
         'port': int(getenv('MLFLOW_PORT', '5080')),
         'username': getenv('MLFLOW_USERNAME', 'aignosi'),
-        'password': getenv('MLFLOW_PASSWORD', 'aignosi')
+        'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
     }
 
 
-def build_opc_config() -> Dict[str, Any]:
+def build_opc_config() -> dict[str, Any]:
     """
     Build OPC server configuration from environment variables.
 
@@ -75,7 +75,7 @@ def build_opc_config() -> Dict[str, Any]:
         OPC_CERT_PATH: Client certificate path (fallback, default: None)
         OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None)
         OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None)
-        OPC_RECONNECTION_INTERVAL: Reconnection interval in milliseconds (fallback, default: 120)
+        OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120)
 
     Returns:
         dict: OPC server configuration dictionary
@@ -93,12 +93,12 @@ def build_opc_config() -> Dict[str, Any]:
             'cert_path': getenv('OPC_CERT_PATH', None),
             'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
             'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
-            'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120'))
+            'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')),
         }
     }
 
 
-def build_mongodb_config() -> Dict[str, Any]:
+def build_mongodb_config() -> dict[str, Any]:
     """
     Build MongoDB configuration from environment variables.
 
@@ -125,5 +125,28 @@ def build_mongodb_config() -> Dict[str, Any]:
     return {
         'connection_string': connection_string,
         'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
-        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600
+        'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
+    }
+
+
+def build_minio_config() -> dict[str, Any]:
+    """
+    Build MinIO (S3-compatible) configuration from environment variables.
+
+    Environment Variables:
+        MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
+        MINIO_ACCESS_KEY: Access key (default: minioadmin)
+        MINIO_SECRET_KEY: Secret key (default: minioadmin)
+        MINIO_REGION: Region name for S3 client (default: us-east-1)
+        MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
+
+    Returns:
+        dict: MinIO configuration dictionary
+    """
+    return {
+        'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
+        'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
+        'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
+        'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'),
+        'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'),
     }

laborious/utils/filters/conditional_filters.py

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

laborious/utils/filters/mlflow_filters.py

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

laborious/utils/repository/minio_repository.py

@@ -0,0 +1,146 @@
+"""
+MinIO repository utilities.
+
+This module provides a lightweight repository around a MinIO/S3-compatible
+object storage using boto3. It supports creating buckets on demand and
+storing/loading pandas DataFrames in Parquet format.
+"""
+
+from io import BytesIO
+from typing import Any
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import ClientError
+from pandas import DataFrame, read_parquet
+from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
+from sientia_do.observability.logger import Logger
+
+
+class MinioRepository:
+    """
+    Repository for interacting with a MinIO (S3-compatible) object storage.
+
+    This class encapsulates a reusable `boto3` S3 client and convenience
+    helpers to persist and retrieve pandas DataFrames as Parquet files.
+
+    Attributes:
+        storage_options (dict): Options compatible with pandas s3fs usage.
+        minio_bucket (str): Default bucket name used for operations.
+        minio_endpoint_url (str): MinIO endpoint URL.
+        minio_region_name (str): MinIO region name.
+        s3_client (Any): Reusable S3 client from `boto3`.
+        logger (Logger): Observability logger.
+        notification_handler (NotificationHandler): Notifications handler.
+    """
+    def __init__(
+        self,
+        minio_endpoint_url: str,
+        minio_access_key: str,
+        minio_secret_key: str,
+        minio_region_name: str,
+        minio_default_bucket: str,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+    ):
+        """Initialize the repository and S3 client.
+
+        Args:
+            minio_endpoint_url (str): MinIO endpoint URL.
+            minio_access_key (str): Access key (AK).
+            minio_secret_key (str): Secret key (SK).
+            minio_region_name (str): Region name for the client.
+            minio_default_bucket (str): Default bucket name to operate on.
+            logger (Logger): Logger instance for structured logs.
+            notification_handler (NotificationHandler): Notification handler.
+        """
+        # MinIO settings shared with pandas s3fs
+        self.storage_options = {
+            'key': minio_access_key,
+            'secret': minio_secret_key,
+            'client_kwargs': {'endpoint_url': minio_endpoint_url},
+        }
+        self.minio_bucket = minio_default_bucket
+        self.minio_endpoint_url = minio_endpoint_url
+        self.minio_region_name = minio_region_name
+
+        logger.info(
+            f'Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}'
+        )
+
+        # Reusable MinIO client
+        self.s3_client: Any = boto3.client(
+            's3',
+            endpoint_url=self.minio_endpoint_url,
+            aws_access_key_id=self.storage_options['key'],
+            aws_secret_access_key=self.storage_options['secret'],
+            region_name=self.minio_region_name,
+            config=Config(
+                signature_version='s3v4',
+                s3={'addressing_style': 'path'},
+                retries={'max_attempts': 5, 'mode': 'standard'},
+                connect_timeout=5,
+                read_timeout=120,
+            ),
+        )
+
+        self.logger = logger
+        self.notification_handler = notification_handler
+
+    def close(self):
+        """Close the underlying S3 client."""
+        self.s3_client.close()
+
+    def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None:
+        """Ensure the default bucket exists; create it if missing.
+
+        Args:
+            metadata (dict[str, Any]): Metadata used for structured logging.
+        """
+
+        try:
+            self.logger.custom_info(f"Checking if bucket '{self.minio_bucket}' exists", metadata)
+            self.s3_client.head_bucket(Bucket=self.minio_bucket)
+        except ClientError:
+            self.logger.custom_info(f"Creating bucket '{self.minio_bucket}'", metadata)
+            self.s3_client.create_bucket(Bucket=self.minio_bucket)
+
+    def store_dataframe_as_parquet(
+        self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any]
+    ):
+        """Persist a DataFrame as a Parquet object in the default bucket.
+
+        Args:
+            dataframe (DataFrame): DataFrame to persist.
+            uri (str): Human-friendly URI used for logging context.
+            object_name (str): Object key (path/key within the bucket).
+            metadata (dict[str, Any]): Metadata used for structured logging.
+        """
+        self.ensure_bucket_exists(metadata)
+
+        self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata)
+
+        buffer = BytesIO()
+        dataframe.to_parquet(buffer, engine='pyarrow', index=True)
+        buffer.seek(0)
+        self.s3_client.put_object(Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue())
+
+        self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata)
+
+    def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame:
+        """Load a Parquet object from the default bucket into a DataFrame.
+
+        Args:
+            object_key (str): Object key to retrieve from the bucket.
+            metadata (dict[str, Any]): Metadata used for structured logging.
+
+        Returns:
+            DataFrame: Loaded DataFrame.
+        """
+        self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata)
+
+        response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)
+
+        # Read the content into a BytesIO buffer to support seek operations
+        buffer = BytesIO(response['Body'].read())
+        return read_parquet(buffer)

laborious/utils/repository/opc_repository.py

@@ -1,16 +1,16 @@
-import asyncio
-import traceback
 import time
+import traceback
 from datetime import datetime
 from pathlib import Path
 from typing import Any
+
 from asyncua import Client
 from asyncua.crypto.security_policies import SecurityPolicyBasic256
-from asyncua.ua import DataValue, Variant, VariantType, DateTime
+from asyncua.ua import DataValue, DateTime, Variant, VariantType
-from regex import F
 from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
 from sientia_do.notifications.models import NotificationLevel
 from sientia_do.observability.logger import Logger
+
 from laborious import metrics
 
 data_type_map = {
@@ -33,17 +33,26 @@ data_type_map = {
     'str': {
         'converter': str,
         'opc_type': VariantType.String,
-    }
+    },
 }
 
 
-class OpcRepository():
+class OpcRepository:
-    def __init__(self, id: str, url: str, logger: Logger,
+    def __init__(
+        self,
+        opc_id: str,
+        url: str,
+        logger: Logger,
         notification_handler: NotificationHandler,
-                 reconnection_interval: int = 60, server_uri: str = None, cert_path: str = None,
+        reconnection_interval: int = 60,
-                 private_key_path: str = None, server_cert_path: str = None, pod_id: str = None):
+        server_uri: str | None = None,
+        cert_path: str | None = None,
+        private_key_path: str | None = None,
+        server_cert_path: str | None = None,
+        pod_id: str | None = None,
+    ):
         self.url = url
-        self.id = id
+        self.id = opc_id
         self.server_uri = server_uri
         self.cert_path = cert_path
         self.private_key_path = private_key_path
@@ -51,16 +60,16 @@ class OpcRepository():
         self.logger = logger
         self.error_count = 0
         self.reconnection_interval = reconnection_interval
-        self.last_reconnection_time = None
+        self.last_reconnection_time: None | datetime = None
         self.notification_handler = notification_handler
-        self.client = None
+        self.client: None | Client = None
         self.pod_id = pod_id
 
         self.metadata = {
             'model_name': '-',
             'model_id': '-',
             'workflow_name': 'opc_repository',
-            'schedule_name': '-'
+            'schedule_name': '-',
         }
 
     async def set_security(self):
@@ -83,13 +92,17 @@ class OpcRepository():
             - Session Timeout: 10,000,000 ms
         """
 
-        if not all([self.cert_path, self.private_key_path]):
+        if self.cert_path is None or self.private_key_path is None:
             raise ValueError(
-                "Certificate and private key paths must be provided for secure connection.")
+                'Certificate and private key paths must be provided for secure connection.'
+            )
+
         cert = Path(self.cert_path)
         private_key = Path(self.private_key_path)
-        server_cert = Path(
+        server_cert = Path(self.server_cert_path) if self.server_cert_path else None
-            self.server_cert_path) if self.server_cert_path else None
+
+        if self.client is None:
+            raise ValueError('Client must be initialized before setting security')
 
         self.client.application_uri = self.server_uri
         self.logger.custom_info('Setting security...', self.metadata)
@@ -97,7 +110,7 @@ class OpcRepository():
             SecurityPolicyBasic256,
             certificate=str(cert),
             private_key=str(private_key),
-            server_certificate=str(server_cert)
+            server_certificate=str(server_cert) if server_cert else None,
         )
         self.client.secure_channel_timeout = 10000000
         self.client.session_timeout = 10000000
@@ -115,8 +128,7 @@ class OpcRepository():
         self.client = Client(self.url)
         if self.cert_path:
             await self.set_security()
-        self.logger.custom_info(
+        self.logger.custom_info(f'Starting connection to OPC server {self.id}...', self.metadata)
-            f'Starting connection to OPC server {self.id}...', self.metadata)
         return await self.try_connect()
 
     async def try_connect(self) -> tuple[bool, dict[str, Any]]:
@@ -136,6 +148,13 @@ class OpcRepository():
 
         try:
             self.last_reconnection_time = datetime.now()
+            if self.client is None:
+                return False, {
+                    'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
+                    'message': 'Client is not initialized',
+                    'block': 'opc_repository',
+                    'level': NotificationLevel.ERROR,
+                }
             await self.client.connect()
             return True, {}
         except Exception as e:
@@ -143,11 +162,11 @@ class OpcRepository():
             self.logger.custom_error(trace, self.metadata)
 
             return False, {
-                "notification_id": f"OPC_CONNECTION_ERROR_{self.id}",
+                'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
-                "message": f"Failed to connect to OPC server: {e}",
+                'message': f'Failed to connect to OPC server: {e}',
-                "block": "opc_repository",
+                'block': 'opc_repository',
-                "level": NotificationLevel.ERROR,
+                'level': NotificationLevel.ERROR,
-                "attachment_content": trace
+                'attachment_content': trace,
             }
 
     async def disconnect(self):
@@ -162,11 +181,9 @@ class OpcRepository():
             return
         try:
             await self.client.disconnect()
-            self.logger.custom_info(
+            self.logger.custom_info('Disconnected from OPC server', self.metadata)
-                'Disconnected from OPC server', self.metadata)
         except Exception as e:
-            self.logger.custom_error(
+            self.logger.custom_error(f'Failed to disconnect from OPC server: {e}', self.metadata)
-                f"Failed to disconnect from OPC server: {e}", self.metadata)
         self.client = None
 
     async def validate_connection(self) -> tuple[bool, dict[str, Any]]:
@@ -203,52 +220,62 @@ class OpcRepository():
 
         if self.error_count > 5:
             self.logger.custom_warning(
-                f"OPC server {self.id} will be disconnected due to multiple errors", self.metadata)
+                f'OPC server {self.id} will be disconnected due to multiple errors', self.metadata
+            )
             try:
                 await self.disconnect()
             except Exception as e:
                 trace = traceback.format_exc()
                 self.logger.custom_error(
-                    f"Failed to disconnect from OPC server: {e}", self.metadata)
+                    f'Failed to disconnect from OPC server: {e}', self.metadata
+                )
                 self.logger.custom_error(trace, self.metadata)
             self.logger.custom_info(
-                f"Attempting to reconnect to OPC server {self.id}...", self.metadata)
+                f'Attempting to reconnect to OPC server {self.id}...', self.metadata
+            )
             return await self.connect()
 
         # Check if client is connected using asyncua's connection state
         try:
-            if self.client.uaclient.protocol is None or self.client.uaclient.protocol.state == "closed":
+            if (
+                self.client.uaclient.protocol is None
+                or self.client.uaclient.protocol.state == 'closed'
+            ):
                 # OPC server is not connected
-                self.logger.custom_error(
+                self.logger.custom_error(f'OPC server {self.id} is not connected', self.metadata)
-                    f"OPC server {self.id} is not connected", self.metadata)
+                if (
-                if self.last_reconnection_time is None or (datetime.now() - self.last_reconnection_time).total_seconds(
+                    self.last_reconnection_time is None
-                ) > self.reconnection_interval:
+                    or (datetime.now() - self.last_reconnection_time).total_seconds()
+                    > self.reconnection_interval
+                ):
                     await self.disconnect()
                     self.logger.custom_info(
-                        f"Trying to reconnect to OPC server {self.id}...", self.metadata)
+                        f'Trying to reconnect to OPC server {self.id}...', self.metadata
+                    )
                     return await self.connect()
 
                 return False, {
-                    "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{self.id}",
+                    'notification_id': f'OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{self.id}',
-                    "message": f"OPC server {self.id} is not connected, waiting for next reconnection window...",
+                    'message': f'OPC server {self.id} is not connected, waiting for next reconnection window...',
-                    "block": "opc_repository",
+                    'block': 'opc_repository',
-                    "level": NotificationLevel.WARNING
+                    'level': NotificationLevel.WARNING,
                 }
             return True, {}
         except Exception as e:
             trace = traceback.format_exc()
-            message = f"Failed to validate connection to OPC server: {e}"
+            message = f'Failed to validate connection to OPC server: {e}'
             self.logger.custom_error(message, self.metadata)
             return False, {
-                "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{self.id}",
+                'notification_id': f'OPC_CONNECTION_CHECK_ERROR_{self.id}',
-                "message": message,
+                'message': message,
-                "block": "opc_repository",
+                'block': 'opc_repository',
-                "level": NotificationLevel.ERROR,
+                'level': NotificationLevel.ERROR,
-                "attachment_content": trace
+                'attachment_content': trace,
             }
 
-    async def write_data(self, node: str, value: Any, data_type: str,
+    async def write_data(
-                         logger: Logger, metadata: dict[str, Any]) -> tuple[bool, dict[str, Any]]:
+        self, node: str, value: Any, data_type: str, logger: Logger, metadata: dict[str, Any]
+    ) -> tuple[bool, dict[str, Any]]:
         """
         Write data to OPC server with comprehensive validation and monitoring.
 
@@ -286,42 +313,36 @@ class OpcRepository():
         start_time = time.time()
 
         try:
-            node_obj = self.client.get_node(node)
+            # ignored because self.validate_connection is called before, so we know self.client is not None
+            node_obj = self.client.get_node(node)  # type: ignore[union-attr]
         except Exception as e:
             trace = traceback.format_exc()
             logger.custom_error(trace, metadata.get('schedule_name', 'N/A'))
             self.error_count += 1
             return False, {
-                "notification_id": f"OPC_WRITE_GET_NODE_ERROR_{self.id}",
+                'notification_id': f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
-                "message": f"Failed to get node from OPC server: {e} | metadata: {metadata}",
+                'message': f'Failed to get node from OPC server: {e} | metadata: {metadata}',
-                "block": "opc_repository",
+                'block': 'opc_repository',
-                "level": NotificationLevel.ERROR,
+                'level': NotificationLevel.ERROR,
-                "attachment_content": trace
+                'attachment_content': trace,
             }
 
         if data_type not in data_type_map:
             return False, {
-                "notification_id": f"OPC_WRITE_DATA_TYPE_ERROR_{self.id}",
+                'notification_id': f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}',
-                "message": f"Unsupported data type: {data_type} | metadata: {metadata}",
+                'message': f'Unsupported data type: {data_type} | metadata: {metadata}',
-                "block": "opc_repository",
+                'block': 'opc_repository',
-                "level": NotificationLevel.ERROR
+                'level': NotificationLevel.ERROR,
             }
 
         data = data_type_map[data_type]['converter'](value)
-        logger.custom_info(
+        logger.custom_info(f'Writing {data} - {type(data)} to {node}', metadata)
-            f'Writing {data} - {type(data)} to {node}', metadata)
         now = datetime.now()
         ua_data = DataValue(
             Variant(data, data_type_map[data_type]['opc_type']),
             SourceTimestamp=DateTime(
-                now.year,
+                now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond
-                now.month,
+            ),
-                now.day,
-                now.hour,
-                now.minute,
-                now.second,
-                now.microsecond
-            )
         )
 
         try:
@@ -331,7 +352,7 @@ class OpcRepository():
                 pod_id=self.pod_id,
                 model_name=metadata['model_name'],
                 pipeline_name=metadata['workflow_name'],
-                opc_server_id=self.id
+                opc_server_id=self.id,
             ).inc()
 
             end_time = time.time()
@@ -340,7 +361,7 @@ class OpcRepository():
                 pod_id=self.pod_id,
                 model_name=metadata['model_name'],
                 pipeline_name=metadata['workflow_name'],
-                opc_server_id=self.id
+                opc_server_id=self.id,
             ).observe(response_time)
 
         except Exception as e:
@@ -348,11 +369,11 @@ class OpcRepository():
             logger.custom_error(trace, metadata)
             self.error_count += 1
             return False, {
-                "notification_id": f"OPC_WRITE_DATA_ERROR_{self.id}",
+                'notification_id': f'OPC_WRITE_DATA_ERROR_{self.id}',
-                "message": f"Failed to write data to OPC server: {e} | metadata: {metadata}",
+                'message': f'Failed to write data to OPC server: {e} | metadata: {metadata}',
-                "block": "opc_repository",
+                'block': 'opc_repository',
-                "level": NotificationLevel.ERROR,
+                'level': NotificationLevel.ERROR,
-                "attachment_content": trace
+                'attachment_content': trace,
             }
         self.error_count = 0
 

laborious/worker/worker.py

@@ -25,33 +25,37 @@ Environment Variables:
 - PROJECT_NAME: Project name for notifications (default: laborious)
 """
 
-from temporalio import workflow, client
+from temporalio import client, workflow
-from temporalio.worker import Worker, PollerBehaviorAutoscaling
+from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
-from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
+from temporalio.worker import PollerBehaviorAutoscaling, Worker
 
 with workflow.unsafe.imports_passed_through():
+    import asyncio
     import os
     import sys
-    import asyncio
+
-    from laborious.workflows.minimal_retrain import MinimalRetrain
+    from prometheus_client import start_http_server
-    from laborious.workflows.predictions_batch import PredictionsBatch
-    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
-    from laborious.workflows.sub_workflows.format_and_export_prediction import \
-        FormatAndExportPrediction
-    from laborious.activities.activities import Activities
-    from laborious.utils.connectors_config import (
-        build_postgres_config,
-        build_mlflow_config,
-        build_opc_config,
-        build_mongodb_config
-    )
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.observability.logger import get_logger
+
     from laborious import metrics
-    from prometheus_client import start_http_server
+    from laborious.activities.activities import Activities
+    from laborious.utils.connectors_config import (
+        build_minio_config,
+        build_mlflow_config,
+        build_mongodb_config,
+        build_opc_config,
+        build_postgres_config,
+    )
+    from laborious.workflows.minimal_retrain import MinimalRetrain
+    from laborious.workflows.predictions_batch import PredictionsBatch
+    from laborious.workflows.sub_workflows.format_and_export_prediction import (
+        FormatAndExportPrediction,
+    )
+    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
 
 POD_ID = os.getenv('POD_ID')
-SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
+SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
 
 
 async def main():
@@ -87,7 +91,7 @@ async def main():
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
 
-    logger.custom_info("Starting prometheus client...", metadata)
+    logger.custom_info('Starting prometheus client...', metadata)
     start_prometheus_server()
 
     logger.custom_info('Starting Notification Handler...', metadata)
@@ -97,7 +101,7 @@ async def main():
         connection_string=mongo_config['connection_string'],
         database=mongo_config['database_name'],
         logger=logger,
-        project_name=os.getenv('PROJECT_NAME', 'laborious')
+        project_name=os.getenv('PROJECT_NAME', 'laborious'),
     )
 
     logger.custom_info('Starting Activities...', metadata)
@@ -105,21 +109,20 @@ async def main():
     activities = Activities(
         postgres_config=build_postgres_config(),
         mlflow_config=build_mlflow_config(),
+        minio_config=build_minio_config(),
         opc_config=build_opc_config(),
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     logger.custom_info('Initializing OPC...', metadata)
     await activities.init_opc()
 
-    logger.custom_info(
+    logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
-        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
 
     new_runtime = Runtime(
         telemetry=TelemetryConfig(
-            metrics=PrometheusConfig(
+            metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
-                bind_address=f"0.0.0.0:{SDK_METRICS_PORT}")
         )
     )
 
@@ -128,7 +131,7 @@ async def main():
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
-        runtime=new_runtime
+        runtime=new_runtime,
     )
 
     logger.custom_info('Starting Workers...', metadata)
@@ -140,26 +143,28 @@ async def main():
             workflows=[MinimalRetrain],
             activities=[
                 activities.load_custom_query,
+                activities.query_to_minio,
                 activities.retrain_model,
                 activities.update_production_model,
-                activities.export_data_to_postgres
+                activities.format_retrain_report,
+                activities.export_data_to_postgres,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,
             max_concurrent_local_activities=50,
             max_cached_workflows=200,
             workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
+            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
         ),
         Worker(
             temporal_client,
             task_queue='predictions_batch-queue',
-            workflows=[PredictionsBatch, PredictionProcess,
+            workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
-                       FormatAndExportPrediction],
             activities=[
                 # MLFlow
                 activities.request_predict,
                 activities.request_transform,
+                activities.query_to_minio,
                 # Gates
                 activities.input_gate,
                 activities.mlflow_response_gate,
@@ -173,15 +178,15 @@ async def main():
                 activities.load_custom_query,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres,
-                activities.write_metrics
+                activities.write_metrics,
             ],
             max_concurrent_workflow_tasks=50,
             max_concurrent_activities=50,
             max_concurrent_local_activities=50,
             max_cached_workflows=200,
             workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
-            activity_task_poller_behavior=PollerBehaviorAutoscaling()
+            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
-        )
+        ),
     ]
 
     handlers = []
@@ -195,7 +200,7 @@ async def main():
         # If an exception occurs in any of the worker handlers, it will be propagated here.
         await asyncio.gather(*handlers)
     except BaseException as e:  # NOSONAR
-        logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
+        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
     finally:
         if notification_handler:
             notification_handler.shutdown()
@@ -224,12 +229,12 @@ def start_prometheus_server():
         SystemExit: If the metrics server fails to start
     """
     try:
-        port = int(os.getenv("HTTP_METRICS_PORT", 9090))
+        port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f"Prometheus server started on port {port}.")
+        print(f'Prometheus server started on port {port}.')
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:
-        print(f"Failed to start Prometheus server: {e}")
+        print(f'Failed to start Prometheus server: {e}')
         os._exit(1)
 
 

laborious/workflows/minimal_retrain.py

@@ -1,14 +1,16 @@
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
-    from laborious.activities.activities import Activities
-    from typing import Any
-    from sientia_do.temporal.policies import retry_policy
     from datetime import timedelta
+    from typing import Any
+
+    from sientia_do.temporal.policies import retry_policy
+
+    from laborious.activities.activities import Activities
 
 
-@workflow.defn(name="minimal_retrain")
+@workflow.defn(name='minimal_retrain')
-class MinimalRetrain():
+class MinimalRetrain:
     """
     Automated model retraining workflow for the Laborious system.
 
@@ -63,44 +65,62 @@ class MinimalRetrain():
                 'schedule_name': input_data['schedule_name'],
                 'model_name': input_data['model_name'],
                 'model_id': input_data['model_id'],
-                'workflow_name': 'minimal_retrain'
+                'workflow_name': 'minimal_retrain',
             }
         }
 
         model_name = input_data['model_name']
+        model_config = input_data.get('model_config', {})
 
-        data = await workflow.execute_local_activity_method(
+        storage_result = await workflow.execute_activity_method(
-            Activities.load_custom_query,
+            Activities.query_to_minio,
             {
                 **metadata,
                 'query': input_data['query'],
-                'datetime_columns': input_data.get('datetime_columns', [])
+                'datetime_columns': input_data.get('datetime_columns', []),
+                'model_name': model_name,
+                'object_prefix': f'retrain_datasets/{model_name}/data',
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600),
         )
 
+        if not storage_result['success']:
+            return
+
         experiment_response = await workflow.execute_activity_method(
             Activities.retrain_model,
             {
                 **metadata,
-                'data': data,
+                'object_key': storage_result['object_key'],
-                'model_name': model_name
+                'model_name': model_name,
+                'model_config': model_config,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(hours=1),
         )
 
-        report = await workflow.execute_activity_method(
+        if experiment_response['success']:
+            update_report = await workflow.execute_activity_method(
                 Activities.update_production_model,
+                {**metadata, 'model_name': model_name, **experiment_response},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(seconds=60),
+            )
+        else:
+            update_report = {}
+
+        report = await workflow.execute_local_activity_method(
+            Activities.format_retrain_report,
             {
                 **metadata,
+                'experiment_response': experiment_response,
                 'model_name': model_name,
                 'model_id': input_data['model_id'],
-                **experiment_response
+                'update_report': update_report,
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=60),
         )
 
         await workflow.execute_activity_method(
@@ -109,8 +129,8 @@ class MinimalRetrain():
                 **metadata,
                 'data': report,
                 'schema': input_data['schema'],
-                'table_name': input_data['table_name']
+                'table_name': input_data['table_name'],
             },
             retry_policy=retry_policy,
-            start_to_close_timeout=timedelta(seconds=60)
+            start_to_close_timeout=timedelta(seconds=600),
         )

laborious/workflows/predictions_batch.py

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

laborious/workflows/sub_workflows/format_and_export_prediction.py

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

laborious/workflows/sub_workflows/prediction_process.py

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

model_convert.ipynb

@@ -0,0 +1,106 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "e838ff21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import csv\n",
+    "\n",
+    "def csv_to_tag_lists(csv_path: str) -> dict:\n",
+    "    read_tags = []\n",
+    "    write_tags = []\n",
+    "\n",
+    "    def to_float(val):\n",
+    "        try:\n",
+    "            return float(str(val).strip())\n",
+    "        except Exception:\n",
+    "            return None\n",
+    "\n",
+    "    with open(csv_path, newline=\"\", encoding=\"utf-8\") as f:\n",
+    "        reader = csv.DictReader(f)\n",
+    "        for row in reader:\n",
+    "            # Basic normalization\n",
+    "            op = (row.get(\"operation\") or \"\").strip()\n",
+    "\n",
+    "            if op == \"READ\":\n",
+    "                # Build common tag payload with required mappings\n",
+    "                tag = {\n",
+    "                    \"server_id\": \"1\",\n",
+    "                    \"tag_address\": row.get(\"opc_tag\"),\n",
+    "                    \"tag_name\": row.get(\"name\"),\n",
+    "                    \"data_range\": [to_float(row.get(\"min_value\")), to_float(row.get(\"max_value\"))],\n",
+    "                    \"aggr_func\": row.get(\"aggregation_func\").lower(),\n",
+    "                    # keep other fields with their original names\n",
+    "                    \"frequency\": row.get(\"frequency\"),\n",
+    "                    \"local\": row.get(\"local\"),\n",
+    "                    \"area\": row.get(\"area\"),\n",
+    "                    \"description\": row.get(\"description\"),\n",
+    "                }\n",
+    "\n",
+    "                read_tags.append(tag)\n",
+    "\n",
+    "            else:\n",
+    "                tag = {\n",
+    "                    \"server_id\": \"1\",\n",
+    "                    \"addr\": row.get(\"opc_tag\"),\n",
+    "                    \"tag_name\": row.get(\"name\"),\n",
+    "                    \"local\": row.get(\"local\"),\n",
+    "                    \"area\": row.get(\"area\"),\n",
+    "                    \"description\": row.get(\"description\"),\n",
+    "                }\n",
+    "                \n",
+    "            if op == \"WRITE_PREDICTION\":\n",
+    "                tag[\"type\"] = \"prediction\"\n",
+    "                write_tags.append(tag)\n",
+    "            elif op == \"WRITE_CONFIDENCE\":\n",
+    "                tag[\"type\"] = \"confidence\"\n",
+    "                write_tags.append(tag)\n",
+    "            # ignore any other operation values silently\n",
+    "\n",
+    "    return {\"read_tags\": read_tags, \"write_tags\": write_tags}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4621cd43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "file_names = [\"Courier - Página1.csv\"]\n",
+    "\n",
+    "for file_name in file_names:\n",
+    "    write_file = file_name.replace(\".csv\", \".json\")\n",
+    "\n",
+    "    with open(write_file, \"w\", encoding=\"utf-8\") as f:\n",
+    "        json.dump(csv_to_tag_lists(file_name), f, indent=2, ensure_ascii=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -0,0 +1,159 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "laborious"
+version = "0.0.0"
+description = "Sientia DataOps Laborious - ML Model Orchestration System"
+readme = "README.md"
+requires-python = ">=3.11"
+authors = [
+    {name = "Aignosi", email = "dev@aignosi.com"}
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+exclude = [
+    ".git",
+    ".venv",
+    "venv",
+    "__pycache__",
+    "*.pyc",
+    ".pytest_cache",
+    "htmlcov",
+]
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle errors
+    "W",   # pycodestyle warnings
+    "F",   # pyflakes
+    "I",   # isort
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "UP",  # pyupgrade
+    "N",   # pep8-naming
+    "YTT", # flake8-2020
+    "S",   # flake8-bandit
+    "BLE", # flake8-blind-except
+    "A",   # flake8-builtins
+    "C90", # mccabe complexity
+]
+
+ignore = [
+    "BLE001",  # ignore blind except, we need to send notifications with any error
+    "E501",  # line too long (handled by formatter)
+    "S101",  # use of assert (needed for tests)
+    "S105",  # possible hardcoded password (false positives)
+    "S106",  # possible hardcoded password (false positives)
+    "N802",  # function name should be lowercase (temporal decorators)
+    "N806",  # variable in function should be lowercase
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = [
+    "S101",  # assert allowed in tests
+    "S105",  # hardcoded passwords ok in tests
+    "S106",  # hardcoded passwords ok in tests
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 15
+
+[tool.ruff.format]
+quote-style = "single"
+indent-style = "space"
+line-ending = "auto"
+
+[tool.mypy]
+python_version = "3.11"
+warn_return_any = false
+warn_unused_configs = true
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+check_untyped_defs = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = false
+warn_no_return = true
+strict_equality = true
+ignore_missing_imports = true
+
+# Ignore missing imports for external packages
+[[tool.mypy.overrides]]
+module = "temporalio.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "sientia_do.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "mlflow.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "prometheus_client.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "sientia.*"
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = "pandas.*"
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+    "-v",
+    "--strict-markers",
+    "--cov=model_manager",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "--cov-report=xml",
+]
+markers = [
+    "asyncio: marks tests as async",
+    "integration: marks tests as integration tests",
+    "unit: marks tests as unit tests",
+]
+
+[tool.coverage.run]
+source = ["model_manager"]
+omit = [
+    "*/tests/*",
+    "*/venv/*",
+    "*/__pycache__/*",
+    "*/site-packages/*",
+]
+branch = true
+
+[tool.coverage.report]
+precision = 2
+show_missing = true
+skip_covered = false
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "def __str__",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+    "class .*\\bProtocol\\):",
+    "@(abc\\.)?abstractmethod",
+]
+
+[tool.coverage.html]
+directory = "htmlcov"
+
+[tool.bandit]
+exclude_dirs = ["tests", "venv", ".venv"]
+skips = ["B101", "B601"]  # Skip assert and shell injection in controlled environments

requirements-dev.txt

@@ -0,0 +1,19 @@
+# Development and Testing Dependencies
+# These packages are only needed for development, testing, and code quality checks
+# Install with: pip install -r requirements-dev.txt
+
+# Code Quality & Linting
+ruff>=0.1.0              # Fast Python linter and formatter (replaces flake8, black, isort)
+mypy>=1.7.0              # Static type checker
+bandit>=1.7.5            # Security vulnerability scanner
+pandas-stubs>=2.0.0      # Type stubs for pandas
+types-requests>=2.31.0   # Type stubs for requests
+
+# Testing
+pytest>=7.4.0            # Testing framework
+pytest-cov>=4.1.0        # Coverage plugin for pytest
+pytest-asyncio>=0.21.0   # Async test support (already in main requirements)
+
+# Development Tools
+ipython>=8.12.0          # Enhanced Python shell
+ipdb>=0.13.13            # IPython debugger

requirements-light.txt

@@ -0,0 +1,12 @@
+temporalio
+psycopg2-binary
+sqlalchemy
+asyncua
+redis
+git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.6
+prometheus-client
+botocore
+boto3
+s3fs
+pyarrow
+mlflow

requirements.txt

@@ -6,3 +6,7 @@ redis
 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.6
 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.39.0
 prometheus-client
+botocore
+boto3
+s3fs
+pyarrow

simulator/Dockerfile

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

tests/laborious/activities/test_activities.py

@@ -1,18 +1,19 @@
+from unittest.mock import ANY, MagicMock, patch
+
 from pytest import mark
-from unittest.mock import patch, MagicMock, ANY
+
-from sientia_do.temporal.activities.postgres import Postgres
 from laborious.activities.activities import Activities
-from laborious.activities.mlflow import MLFlow
 from laborious.activities.gates import Gates
+from laborious.activities.mlflow import MLFlow
 from laborious.activities.opc import OPC
+from laborious.activities.storage import Storage
 
 
-@patch('laborious.activities.activities.Postgres.__init__')
+@patch('laborious.activities.activities.Storage.__init__')
 @patch('laborious.activities.activities.MLFlow.__init__')
 @patch('laborious.activities.activities.OPC.__init__')
 @patch('laborious.activities.activities.Gates.__init__')
-def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgres_init):
+def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_storage_init):
-
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -20,20 +21,23 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
+    minio_config = {
-        'host': 'localhost',
+        'endpoint_url': 'localhost:9000',
-        'port': 5000,
+        'access_key': 'minio',
-        'username': 'mlflow',
+        'secret_key': 'minio123',
-        'password': 'mlflow'
+        'region_name': 'us-east-1',
+        'default_bucket': 'test',
     }
 
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
     opc_config = {
         'bootstrap_servers': 'localhost:9092',
         'polling_time': 1000,
-        'group_id': 'test-group'
+        'group_id': 'test-group',
     }
 
     logger = MagicMock()
@@ -42,18 +46,19 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
+        minio_config=minio_config,
         opc_config=opc_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     assert isinstance(activities, Activities)
-    assert isinstance(activities, Postgres)
+    assert isinstance(activities, Storage)
     assert isinstance(activities, MLFlow)
     assert isinstance(activities, OPC)
     assert isinstance(activities, Gates)
 
-    mock_postgres_init.assert_called_once_with(
+    mock_storage_init.assert_called_once_with(
         ANY,
         host=postgres_config['host'],
         port=postgres_config['port'],
@@ -62,8 +67,9 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         dbname=postgres_config['dbname'],
         min_connections=postgres_config['min_connections'],
         max_connections=postgres_config['max_connections'],
+        minio_config=minio_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     mock_mlflow_init.assert_called_once_with(
@@ -72,30 +78,25 @@ def test___init__(mock_gates_init, mock_opc_init, mock_mlflow_init, mock_postgre
         mlflow_port=mlflow_config['port'],
         mlflow_username=mlflow_config['username'],
         mlflow_password=mlflow_config['password'],
+        minio_config=minio_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     mock_opc_init.assert_called_once_with(
-        ANY,
+        ANY, opc_servers=opc_config, logger=logger, notification_handler=notification_handler
-        opc_servers=opc_config,
-        logger=logger,
-        notification_handler=notification_handler
     )
 
     mock_gates_init.assert_called_once_with(
-        ANY,
+        ANY, logger=logger, notification_handler=notification_handler
-        logger=logger,
-        notification_handler=notification_handler
     )
 
 
 @mark.asyncio
-@patch('laborious.activities.activities.Postgres', return_value=MagicMock())
+@patch('laborious.activities.activities.Storage', return_value=MagicMock())
 @patch('laborious.activities.activities.MLFlow', return_value=MagicMock())
 @patch('laborious.activities.activities.OPC', return_value=MagicMock())
-async def test_shutdown(mock_opc_init,
+async def test_shutdown(mock_opc_init, _mock_mlflow_init, mock_storage_init):
-                        _mock_mlflow_init, mock_postgres_init):
     postgres_config = {
         'host': 'localhost',
         'port': 5432,
@@ -103,20 +104,23 @@ async def test_shutdown(mock_opc_init,
         'password': 'postgres',
         'dbname': 'postgres',
         'min_connections': 1,
-        'max_connections': 10
+        'max_connections': 10,
     }
 
-    mlflow_config = {
+    minio_config = {
-        'host': 'localhost',
+        'endpoint_url': 'localhost:9000',
-        'port': 5000,
+        'access_key': 'minio',
-        'username': 'mlflow',
+        'secret_key': 'minio123',
-        'password': 'mlflow'
+        'region_name': 'us-east-1',
+        'default_bucket': 'test',
     }
 
+    mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}
+
     opc_config = {
         'bootstrap_servers': 'localhost:9092',
         'polling_time': 1000,
-        'group_id': 'test-group'
+        'group_id': 'test-group',
     }
 
     logger = MagicMock()
@@ -125,11 +129,12 @@ async def test_shutdown(mock_opc_init,
     activities = Activities(
         postgres_config=postgres_config,
         mlflow_config=mlflow_config,
+        minio_config=minio_config,
         opc_config=opc_config,
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
     )
 
     await activities.shutdown()
     mock_opc_init.shutdown.assert_called_once()
-    mock_postgres_init.close.assert_called_once()
+    mock_storage_init.close.assert_called_once()

tests/laborious/activities/test_gates.py

@@ -1,6 +1,8 @@
-from unittest.mock import MagicMock, ANY, patch
+from unittest.mock import ANY, MagicMock, patch
+
 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
+
 from laborious.activities.gates import Gates
 
 
@@ -20,11 +22,11 @@ def gates_activity():
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
@@ -34,20 +36,18 @@ async def test_input_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
         'data': {'value': [1, 2, 3]},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.error.assert_called_once_with(
-        "Filter INVALID_FILTER not found", metadata['metadata']
+        'Filter INVALID_FILTER not found', metadata['metadata']
     )
 
 
@@ -57,28 +57,27 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac
     # Arrange
     mock_input_filter_functions.__contains__.return_value = True
     mock_input_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
-        },
         'data': {'value': []},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="INTPUT_GATE_ERROR__EMPTY_DATA",
+        notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
         message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
-        block="input_gate",
+        block='input_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -89,14 +88,14 @@ async def test_input_gate_no_filters(gates_activity):
         **metadata,
         'filters': {},
         'data': {'value': [1, 2, 3]},
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -105,18 +104,34 @@ async def test_input_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
-            'EMPTY_DATA': {'policy': 'STOP', 'config': {}}
-        },
         'data': {'value': []},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.input_gate(input_data)
 
     # Assert
-    assert result == ('STOP', -1, "Input data with bad quality")
+    assert result == ('STOP', -1, 'Input data with bad quality')
+    gates_activity.debug.assert_called()
+
+
+@mark.asyncio
+async def test_input_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
+        'data': {'value': [1, 2, 3]},
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    }
+
+    # Act
+    result = await gates_activity.input_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -125,51 +140,49 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
 
 
 @mark.asyncio
 @patch('laborious.activities.gates.mlflow_response_filter_functions')
-async def test_mlflow_response_gate_filter_exception(mock_mlflow_response_filter_functions,
+async def test_mlflow_response_gate_filter_exception(
-                                                     gates_activity):
+    mock_mlflow_response_filter_functions, gates_activity
+):
     # Arrange
     mock_mlflow_response_filter_functions.__contains__.return_value = True
     mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-            'INVALID_FILTER': {'POLICY': 'STOP'}
-        },
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER",
+        notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
         message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
-        block="mlflow_gate",
+        block='mlflow_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -181,14 +194,14 @@ async def test_mlflow_response_gate_no_filters(gates_activity):
         'filters': {},
         'data': {'content': {'message': 'success'}},
         'type': 'test',
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -197,86 +210,101 @@ async def test_mlflow_response_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'API_ERROR': {'policy': 'STOP'}},
-            'API_ERROR': {'policy': 'STOP'}
-        },
         'data': {
             'success': False,
-            'content': {
+            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
-                'message': 'API error occurred',
-                'traceback': 'error trace'
-            }
         },
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_response_gate(input_data)
 
     # Assert
-    assert result == ('STOP', -1, "API error occurred")
+    assert result == ('STOP', -1, 'API error occurred')
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called()
 
 
+@mark.asyncio
+async def test_mlflow_response_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'API_ERROR': {'policy': 'STOP'}},
+        'data': {
+            'success': True,
+            'content': {'message': 'success'},
+        },
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    }
+
+    # Act
+    result = await gates_activity.mlflow_response_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, '')
+    gates_activity.debug.assert_called()
+
+
 @mark.asyncio
 async def test_mlflow_content_gate_invalid_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-            'INVALID_FILTER': {'POLICY': 'STOP'}
+        'data': {
+            'success': True,
+            'content': {'message': 'success'},
         },
-        'data': {'value': [1, 2, 3]},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
 
 
 @mark.asyncio
 @patch('laborious.activities.gates.mlflow_content_filter_functions')
-async def test_mlflow_content_gate_filter_exception(mock_mlflow_content_filter_functions,
+async def test_mlflow_content_gate_filter_exception(
-                                                    gates_activity):
+    mock_mlflow_content_filter_functions, gates_activity
+):
     # Arrange
     mock_mlflow_content_filter_functions.__contains__.return_value = True
     mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
-        side_effect=Exception("Test error"))
+        side_effect=Exception('Test error')
+    )
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'API_ERROR': {'POLICY': 'STOP'}},
-            'API_ERROR': {'POLICY': 'STOP'}
-        },
         'data': {
             'success': False,
-            'content': {
+            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
-                'message': 'API error occurred',
-                'traceback': 'error trace'
-            }
         },
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
-        notification_id="MLFLOW_GATE_CONTENT_FILTER__API_ERROR",
+        notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
         message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
-        block="mlflow_gate",
+        block='mlflow_gate',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
@@ -288,14 +316,14 @@ async def test_mlflow_content_gate_no_filters(gates_activity):
         'filters': {},
         'data': {'value': [1, 2, 3]},
         'type': 'test',
-        'path_priority': ['CONTINUE', 'STOP', 'REPEAT']
+        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (None, 0, "")
+    assert result == (None, 0, '')
     gates_activity.debug.assert_called()
 
 
@@ -304,31 +332,48 @@ async def test_mlflow_content_gate_with_filter(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'filters': {
+        'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
-            'NAN_VALUES': {'policy': 'STOP', 'config': {}}
-        },
         'data': {'value': [None, None, None]},
         'type': 'test',
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT']
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
     # Act
     result = await gates_activity.mlflow_content_gate(input_data)
 
     # Assert
-    assert result == (
+    assert result == ('STOP', -1, 'Transformed data not passed the content filter')
-        'STOP', -1, "Transformed data not passed the content filter")
     gates_activity.debug.assert_called()
     gates_activity.send_notification.assert_called()
 
 
+@mark.asyncio
+async def test_mlflow_content_gate_with_filter_not_caught(gates_activity):
+    # Arrange
+    input_data = {
+        **metadata,
+        'filters': {'API_ERROR': {'POLICY': 'STOP'}},
+        'data': {'content': {'message': 'success'}},
+        'type': 'test',
+        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    }
+
+    # Act
+    result = await gates_activity.mlflow_content_gate(input_data)
+
+    # Assert
+    assert result == (None, 0, '')
+    gates_activity.debug.assert_called()
+
+
 def test_get_prediction_store_policy_invalid_policy(gates_activity):
     # Arrange
     prediction_store_policy = 'INVALID_POLICY'
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -341,7 +386,8 @@ def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -354,7 +400,8 @@ def test_get_prediction_store_policy_valid_policy_type(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'lts'
@@ -367,7 +414,8 @@ def test_get_prediction_store_policy_valid_policy(gates_activity):
 
     # Act
     policy_type, policy_value = gates_activity.get_prediction_store_policy(
-        prediction_store_policy, metadata)
+        prediction_store_policy, metadata
+    )
 
     # Assert
     assert policy_type == 'erl'
@@ -380,16 +428,12 @@ async def test_format_prediction_no_timestamp(gates_activity):
     input_data = {
         **metadata,
         'data': {
-            'prediction': {
+            'prediction': {'2023-05-26 11:12:27': 1},
-                '2023-05-26 11:12:27': 1
+            'response_time': {'2023-05-26 11:12:27': 0.1},
-            },
-            'response_time': {
-                '2023-05-26 11:12:27': 0.1
-            }
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:1'
+        'prediction_store_policy': 'lts:1',
     }
 
     # Act
@@ -402,7 +446,7 @@ async def test_format_prediction_no_timestamp(gates_activity):
     assert result['model_id'] == {0: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9}
     assert result['prediction_status'] == {0: 'Good'}
-    assert result['comments'] == {0: ""}
+    assert result['comments'] == {0: ''}
 
 
 @mark.asyncio
@@ -420,11 +464,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
                 '2023-05-26 11:12:27': 0.1,
                 '2023-05-26 11:12:28': 0.2,
                 '2023-05-26 11:12:29': 0.3,
-            }
+            },
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'erl:2'
+        'prediction_store_policy': 'erl:2',
     }
 
     # Act
@@ -433,12 +477,11 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
     # Assert
     assert result['prediction'] == {0: 2, 1: 1}
     assert result['response_time'] == {0: 0.2, 1: 0.1}
-    assert result['timestamp'] == {
+    assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
-        0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
     assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
     assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
-    assert result['comments'] == {0: "", 1: ""}
+    assert result['comments'] == {0: '', 1: ''}
 
 
 @mark.asyncio
@@ -456,11 +499,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
                 '2023-05-26 11:12:27': 0.1,
                 '2023-05-26 11:12:28': 0.2,
                 '2023-05-26 11:12:29': 0.3,
-            }
+            },
         },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:2'
+        'prediction_store_policy': 'lts:2',
     }
 
     # Act
@@ -469,12 +512,11 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
     # Assert
     assert result['prediction'] == {0: 3, 1: 2}
     assert result['response_time'] == {0: 0.3, 1: 0.2}
-    assert result['timestamp'] == {
+    assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
-        0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
     assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
     assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
     assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
-    assert result['comments'] == {0: "", 1: ""}
+    assert result['comments'] == {0: '', 1: ''}
 
 
 @mark.asyncio
@@ -482,22 +524,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {'prediction': [1, 2, 3],
+        'data': {
+            'prediction': [1, 2, 3],
             'response_time': [0.1, 0.2, 0.3],
-                 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']},
+            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
+        },
         'model_id': 'test_model',
         'prediction_confidence': 0.9,
-        'prediction_store_policy': 'lts:2'
+        'prediction_store_policy': 'lts:2',
     }
-    gates_activity.get_prediction_store_policy = MagicMock(
+    gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
-        return_value=('invalid', 1))
 
     try:
-        result = await gates_activity.format_prediction(input_data)
+        await gates_activity.format_prediction(input_data)
     except ValueError as e:
-        assert str(e) == "Invalid policy type: invalid"
+        assert str(e) == 'Invalid policy type: invalid'
     else:
-        assert False, "Expected ValueError"
+        raise AssertionError('Expected ValueError')
 
 
 @mark.asyncio
@@ -508,7 +551,7 @@ async def test_format_default_prediction(gates_activity):
         'timestamp': '2023-05-26 11:12:27',
         'model_id': 'test_model',
         'prediction_confidence': 0.1,
-        'comment': 'Test comment'
+        'comment': 'Test comment',
     }
 
     # Act
@@ -526,15 +569,42 @@ async def test_format_default_prediction(gates_activity):
 
 
 @mark.asyncio
-async def test_get_last_timestamp_with_data(gates_activity):
+async def test_format_retrain_report(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'experiment_response': {
-            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']
+            'success': True,
-        }
+            'timestamp': '2023-05-26 11:12:27',
+            'message': 'success',
+        },
+        'update_report': {
+            'version': '1.0.0',
+            'mlflow_run_id': 'test_mlflow_run_id',
+            'mlflow_experiment_id': 'test_mlflow_experiment_id',
+        },
+        'model_id': 'test_model',
+        'model_name': 'test_model',
     }
 
+    # Act
+    result = await gates_activity.format_retrain_report(input_data)
+
+    # Assert
+    assert result['model_id'] == {0: 'test_model'}
+    assert result['model_name'] == {0: 'test_model'}
+    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
+    assert result['status'] == {0: 'success'}
+    assert result['version'] == {0: '1.0.0'}
+    assert result['mlflow_run_id'] == {0: 'test_mlflow_run_id'}
+    assert result['mlflow_experiment_id'] == {0: 'test_mlflow_experiment_id'}
+
+
+@mark.asyncio
+async def test_get_last_timestamp_with_data(gates_activity):
+    # Arrange
+    input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}
+
     # Act
     result = await gates_activity.get_last_timestamp(input_data)
 
@@ -545,10 +615,7 @@ async def test_get_last_timestamp_with_data(gates_activity):
 @mark.asyncio
 async def test_get_last_timestamp_no_data(gates_activity):
     # Arrange
-    input_data = {
+    input_data = {'data': {}, **metadata}
-        'data': {},
-        **metadata
-    }
 
     # Act
     result = await gates_activity.get_last_timestamp(input_data)
@@ -567,30 +634,28 @@ async def test_write_metrics(mock_metrics, gates_activity):
         'prediction': {
             'prediction': [1, 2, 3],
             'prediction_confidence': [0.9, 0.8, 0.7],
-            'response_time': [0.1, 0.2, 0.3]
+            'response_time': [0.1, 0.2, 0.3],
-        }
+        },
     }
     await gates_activity.write_metrics(input_data)
     mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
+        pipeline_name=metadata['metadata']['workflow_name'],
     )
     mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()
 
     mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
+        pipeline_name=metadata['metadata']['workflow_name'],
-    )
-    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(
-        0.9
     )
+    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)
 
     mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
         pod_id=gates_activity.pod_id,
         model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name']
+        pipeline_name=metadata['metadata']['workflow_name'],
     )
     mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
         0.1

tests/laborious/activities/test_mlflow.py

@@ -1,45 +1,68 @@
-from datetime import datetime
+from unittest.mock import ANY, MagicMock, call, patch
-from unittest.mock import ANY, MagicMock, patch
 
 import numpy as np
-from pandas import DataFrame, Timestamp
 from pytest import fixture, mark, raises
-from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
-from laborious.activities.mlflow import MLFlow
 from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
+
+from laborious.activities.mlflow import MLFlow
 
 
-@patch("laborious.activities.mlflow.MLFlowRepository")
+@patch('laborious.activities.mlflow.MLFlowRepository')
-def test___init__(mock_mlflow_repository):
+@patch('laborious.activities.mlflow.MinioRepository')
+def test___init__(mock_minio_repository, mock_mlflow_repository):
     mlflow = MLFlow(
-        mlflow_host="http://localhost",
+        mlflow_host='http://localhost',
         mlflow_port=5000,
-        mlflow_username="admin",
+        mlflow_username='admin',
-        mlflow_password="admin",
+        mlflow_password='admin',
+        minio_config={
+            'endpoint_url': 'http://localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
     )
 
-    assert mlflow.mlflow_host == "http://localhost"
+    assert mlflow.mlflow_host == 'http://localhost'
     assert mlflow.mlflow_port == 5000
-    assert mlflow.mlflow_username == "admin"
+    assert mlflow.mlflow_username == 'admin'
-    assert mlflow.mlflow_password == "admin"
+    assert mlflow.mlflow_password == 'admin'
 
-    mock_mlflow_repository.assert_called_once_with(
+    mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
-        "http://localhost:5000", "admin", "admin", ANY
+
+    mock_minio_repository.assert_called_once_with(
+        logger=ANY,
+        notification_handler=ANY,
+        minio_endpoint_url='http://localhost:9000',
+        minio_access_key='minio',
+        minio_secret_key='minio123',
+        minio_region_name='us-east-1',
+        minio_default_bucket='test',
     )
 
 
 @fixture
-@patch("laborious.activities.mlflow.MLFlowRepository")
+@patch('laborious.activities.mlflow.MLFlowRepository')
-def mlflow(mock_mlflow_repository):
+@patch('laborious.activities.mlflow.MinioRepository')
+def mlflow(mock_minio_repository, mock_mlflow_repository):
     mlflow = MLFlow(
-        mlflow_host="http://localhost:5000",
+        mlflow_host='http://localhost:5000',
         mlflow_port=5000,
-        mlflow_username="admin",
+        mlflow_username='admin',
-        mlflow_password="admin",
+        mlflow_password='admin',
+        minio_config={
+            'endpoint_url': 'http://localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
         logger=MagicMock(),
-        notification_handler=MagicMock()
+        notification_handler=MagicMock(),
     )
 
     mlflow.send_notification = MagicMock()
@@ -48,44 +71,67 @@ def mlflow(mock_mlflow_repository):
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("laborious.activities.mlflow.DataFrame")
+@patch('laborious.activities.mlflow.DataFrame')
-@patch("laborious.activities.mlflow.max")
+@patch('laborious.activities.mlflow.max')
 async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
         **metadata,
         'data': [
-            {'timestamp': '2024-01-01', 'variable': 'var1',
+            {
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
+                'timestamp': '2024-01-01',
-            {'timestamp': '2024-01-01', 'variable': 'var2',
+                'variable': 'var1',
-                'value': 2.0, 'created_at': '2024-01-01 12:00:00'},
+                'value': 1.0,
-            {'timestamp': '2024-01-02', 'variable': 'var1',
+                'created_at': '2024-01-01 12:00:00',
-                'value': 3.0, 'created_at': '2024-01-02 12:00:00'},
+            },
-            {'timestamp': '2024-01-02', 'variable': 'var2',
+            {
-                'value': 4.0, 'created_at': '2024-01-02 12:00:00'},
+                'timestamp': '2024-01-01',
-            {'timestamp': '2024-01-02', 'variable': 'var1',
+                'variable': 'var2',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'},
+                'value': 2.0,
-            {'timestamp': '2024-01-02', 'variable': 'var2',
+                'created_at': '2024-01-01 12:00:00',
-                'value': 1.0, 'created_at': '2024-01-01 12:00:00'}
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var1',
+                'value': 3.0,
+                'created_at': '2024-01-02 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var2',
+                'value': 4.0,
+                'created_at': '2024-01-02 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var1',
+                'value': 1.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
+            {
+                'timestamp': '2024-01-02',
+                'variable': 'var2',
+                'value': 1.0,
+                'created_at': '2024-01-01 12:00:00',
+            },
         ],
         'model_name': 'test_model',
-        'model_config': {}
+        'model_config': {},
     }
 
     # Mock the transform response
-    expected_response = {'prediction': [0.5, 0.6], 'timestamp': [
+    expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
-        '2024-01-01', '2024-01-02']}
     mlflow.model_monitoring_repository.transform.return_value = expected_response
 
     mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
@@ -114,30 +160,25 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
 
 
 @mark.asyncio
-@patch("laborious.activities.mlflow.DataFrame")
+@patch('laborious.activities.mlflow.DataFrame')
-@patch("laborious.activities.mlflow.to_datetime")
+@patch('laborious.activities.mlflow.to_datetime')
-@patch("laborious.activities.mlflow.max")
+@patch('laborious.activities.mlflow.max')
 async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
     mock_max.return_value = '2024-01-02'
     # Mock input data
     input_data = {
         **metadata,
         'data': {
-            "variable": {
+            'variable': {
-                "2024-01-01": "var1",
+                '2024-01-01': 'var1',
-                "2024-01-02": "var2",
+                '2024-01-02': 'var2',
-                "2024-01-03": "var1",
+                '2024-01-03': 'var1',
-                "2024-01-04": "var2"
+                '2024-01-04': 'var2',
             },
-            "value": {
+            'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
-                "2024-01-01": 1.0,
-                "2024-01-02": 2.0,
-                "2024-01-03": 3.0,
-                "2024-01-04": 4.0
-            }
         },
         'model_name': 'test_model',
-        'model_config': {}
+        'model_config': {},
     }
 
     # Mock the predict response
@@ -148,8 +189,9 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
     response_data = await mlflow.request_predict(input_data)
 
     mock_dataframe.assert_called_once_with(input_data['data'])
-    mock_dataframe.return_value.replace.assert_called_once_with(
+    mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
-        np.nan, None, inplace=True
+    mock_dataframe.return_value.__setitem__.assert_any_call(
+        'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
     )
     mock_dataframe.return_value.__setitem__.assert_any_call(
         'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
@@ -158,9 +200,12 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
     mock_to_datetime.assert_called_once_with(
         mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
     )
-    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(
+    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
-        DATETIME_FORMAT
+
+    mock_to_datetime.assert_called_once_with(
+        mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
     )
+    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
 
     # Verify the response
     assert response_data == expected_response
@@ -172,98 +217,232 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
 
 
 @mark.asyncio
-async def test_retrain_model(mlflow):
+@patch('laborious.activities.mlflow.to_datetime')
-    data = {
+async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
-        "model_id": [4, 5, 6, 7],
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
-        "created_at": [1, 2, 3, 4],
+        'success': True,
-        "timestamp": [1, 1, 2, 2],
+        'experiment': 'test_experiment',
-        "variable": ["var1", "var2", "var1", "var2"],
+        'message': 'Model retrained successfully.',
-        "value": [1, 2, 3, 4]
     }
 
-    mlflow.model_monitoring_repository.retrain_model.return_value = (
+    response = await mlflow.retrain_model(
-        'Model retrained successfully', 'test')
+        {
-
-    response = await mlflow.retrain_model({
             **metadata,
-        'data': data,
+            'object_key': 'test_object_key',
-        'model_name': 'test_model'
+            'model_name': 'test_model',
-    })
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
 
-    mlflow.model_monitoring_repository.retrain_model.assert_called_once()
+    raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
+    raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
+        subset=['variable', 'timestamp'], keep='first'
+    )
+    raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
 
     assert response == {
-        "status": 'Model retrained successfully',
+        'success': True,
-        "timestamp": 2,
+        'experiment': 'test_experiment',
-        "experiment": 'test'
+        'message': 'Model retrained successfully.',
+        'timestamp': timestamp,
     }
 
 
 @mark.asyncio
-async def test_retrain_model_error(mlflow):
+@patch('laborious.activities.mlflow.to_datetime')
-    mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
+async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
-        'Error retraining model'
+    mlflow.model_monitoring_repository.retrain_model.return_value = {
-    )
+        'success': False,
-
+        'traceback': 'test_traceback',
-    data = {
+        'message': 'Model retrained failed.',
-        "model_id": [4, 5, 6, 7],
-        "created_at": [1, 2, 3, 4],
-        "timestamp": [1, 1, 2, 2],
-        "variable": ["var1", "var2", "var1", "var2"],
-        "value": [1, 2, 3, 4]
     }
 
-    try:
+    response = await mlflow.retrain_model(
-        await mlflow.retrain_model({
+        {
             **metadata,
-            'data': data,
+            'object_key': 'test_object_key',
-            'model_name': 'test_model'
+            'model_name': 'test_model',
-        })
+            'model_config': {
-    except Exception as e:
+                'target': 'target',
-        assert str(e) == 'Error retraining model'
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
+
+    timestamp = raw_data.__getitem__.return_value.max.return_value
+
+    raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
+    raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
+        subset=['variable', 'timestamp'], keep='first'
+    )
+    raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
+
+    raw_data.drop.assert_has_calls(
+        [
+            call(columns=['model_id'], inplace=True, errors='ignore'),
+            call(columns=['created_at'], inplace=True, errors='ignore'),
+        ]
+    )
+    raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
+    raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
+
+    raw_data = raw_data.pivot.return_value
+
+    raw_data.__setitem__.assert_has_calls(
+        [
+            call('timestamp', raw_data.index),
+            call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
+            call('timestamp', mock_to_datetime.return_value),
+        ]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
+    )
+    mock_to_datetime.assert_has_calls(
+        [call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
+    )
+
+    mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
+        data=raw_data,
+        model_name='test_model',
+        model_config={
+            'target': 'target',
+            'transform_flavor': 'sklearn',
+            'predict_flavor': 'pyfunc',
+        },
+        metadata=metadata['metadata'],
+    )
+
     mlflow.send_notification.assert_called_once_with(
         metadata=metadata['metadata'],
         notification_id='RETRAIN_MODEL_ERROR',
-            message='Error retraining model test_model: Error retraining model',
+        message='Error retraining model test_model: Model retrained failed.',
         block='retrain_model',
         level=NotificationLevel.ERROR,
-            attachment_content=ANY
+        attachment_content=ANY,
     )
-    else:
+
-        assert False, "No exception raised"
+    assert response == {
+        'success': False,
+        'traceback': 'test_traceback',
+        'message': 'Model retrained failed.',
+        'timestamp': timestamp,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error(mlflow):
+    mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception(
+        'Error loading retrain data'
+    )
+
+    response = await mlflow.retrain_model(
+        {
+            **metadata,
+            'object_key': 'test_object_key',
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
+
+    assert response == {
+        'success': False,
+        'message': 'Error loading retrain data: Error loading retrain data',
+        'traceback': ANY,
+        'timestamp': ANY,
+    }
+
+
+@mark.asyncio
+async def test_retrain_model_data_error_no_minio_repository(mlflow):
+    mlflow.minio_repository = None
+
+    with raises(ValueError) as e:
+        await mlflow.retrain_model(
+            {
+                **metadata,
+                'object_key': 'test_object_key',
+                'model_name': 'test_model',
+                'model_config': {
+                    'target': 'target',
+                    'transform_flavor': 'sklearn',
+                    'predict_flavor': 'pyfunc',
+                },
+            }
+        )
+
+        assert str(e.value) == 'Minio repository not initialized'
 
 
 @mark.asyncio
 async def test_update_production_model(mlflow):
-    mlflow.model_monitoring_repository.update_production_model.return_value = (
-        {
-            "data1": 1,
-            "data2": 2
-        }
-    )
-
     input_data = {
         **metadata,
         'model_name': 'test_model',
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     response = await mlflow.update_production_model(input_data)
 
     mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
-        experiment='test', model_name='test_model')
+        experiment='test', model_name='test_model', metadata=metadata['metadata']
+    )
 
-    assert response == {
+    assert response == mlflow.model_monitoring_repository.update_production_model.return_value
-        'data1': {0: 1},
-        'data2': {0: 2},
-        'model_id': {0: 1},
-        'model_name': {0: 'test_model'},
-        'timestamp': {0: 2},
-        'status': {0: 'success'}
-    }
 
 
 @mark.asyncio
@@ -278,7 +457,7 @@ async def test_update_production_model_error(mlflow):
         'model_id': 1,
         'experiment': 'test',
         'timestamp': 2,
-        'status': 'success'
+        'status': 'success',
     }
 
     try:
@@ -291,7 +470,7 @@ async def test_update_production_model_error(mlflow):
             message='Error updating production model test_model: Error updating production model',
             block='update_production_model',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
     else:
-        assert False, "No exception raised"
+        raise AssertionError('No exception raised')

tests/laborious/activities/test_opc.py

@@ -1,57 +1,55 @@
-from unittest.mock import patch, MagicMock, ANY, call, AsyncMock
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
-from pandas import DataFrame
+
-from pytest import fixture, mark
 import pytest_asyncio
+from pandas import DataFrame
+from pytest import mark
 from sientia_do.notifications.models import NotificationLevel
 
 from laborious.activities.opc import OPC
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
 
 def test__init__():
-    servers = {
+    servers = {'server1': 'config'}
-        'server1': 'config'
+    opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
-    }
-    opc = OPC(
-        opc_servers=servers,
-        logger=MagicMock(),
-        notification_handler=MagicMock()
-    )
 
     assert opc.opc_servers == servers
     assert opc.opc_repository == {}
 
 
 @mark.asyncio
-@patch("laborious.activities.opc.OpcRepository")
+@patch('laborious.activities.opc.OpcRepository')
-@patch("laborious.activities.opc.OPC.send_notification")
+@patch('laborious.activities.opc.OPC.send_notification')
 async def test_init_opc(mock_send_notification, mock_opc_repository):
     mock_logger = MagicMock()
     server1 = MagicMock(
-        connect=AsyncMock(return_value=(True, {})),
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
-        write_data=AsyncMock(return_value=(True, {}))
     )
     server2 = MagicMock(
-        connect=AsyncMock(return_value=(True, {})),
+        connect=AsyncMock(return_value=(True, {})), write_data=AsyncMock(return_value=(True, {}))
-        write_data=AsyncMock(return_value=(True, {}))
     )
     server3 = MagicMock(
-        connect=AsyncMock(return_value=(False, {
+        connect=AsyncMock(
+            return_value=(
+                False,
+                {
                     'notification_id': 'OPC_CONNECTION_ERROR_server3',
                     'message': 'Failed to connect to OPC server: Test error',
                     'block': 'opc_repository',
                     'level': NotificationLevel.ERROR,
-            'attachment_content': 'Test error'
+                    'attachment_content': 'Test error',
-        })),
+                },
-        write_data=AsyncMock(return_value=(True, {}))
+            )
+        ),
+        write_data=AsyncMock(return_value=(True, {})),
     )
     mock_opc_repository.side_effect = [server1, server2, server3]
     mock_notification_handler = MagicMock()
@@ -82,12 +80,10 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'private_key_path': '',
             'server_cert_path': '',
             'reconnection_interval': 60,
-        }
+        },
     }
     opc = OPC(
-        opc_servers=servers,
+        opc_servers=servers, logger=mock_logger, notification_handler=mock_notification_handler
-        logger=mock_logger,
-        notification_handler=mock_notification_handler
     )
     await opc.init_opc()
 
@@ -97,57 +93,63 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
     assert opc.opc_repository['server1'] == server1
     assert opc.opc_repository['server2'] == server2
 
-    mock_opc_repository.assert_has_calls([
+    mock_opc_repository.assert_has_calls(
+        [
             call(
-            id="server1",
+                opc_id='server1',
-            url="http://localhost:8080",
+                url='http://localhost:8080',
                 logger=mock_logger,
-            server_uri="opc.tcp://localhost:4840",
+                server_uri='opc.tcp://localhost:4840',
-            cert_path="",
+                cert_path='',
-            private_key_path="",
+                private_key_path='',
-            server_cert_path="",
+                server_cert_path='',
                 notification_handler=mock_notification_handler,
                 reconnection_interval=60,
-            pod_id='localhost'
+                pod_id='localhost',
             ),
-    ])
+        ]
-    mock_opc_repository.assert_has_calls([
+    )
+    mock_opc_repository.assert_has_calls(
+        [
             call(
-            id="server2",
+                opc_id='server2',
-            url="http://localhost:8080",
+                url='http://localhost:8080',
                 logger=mock_logger,
-            server_uri="opc.tcp://localhost:4840",
+                server_uri='opc.tcp://localhost:4840',
-            cert_path="",
+                cert_path='',
-            private_key_path="",
+                private_key_path='',
-            server_cert_path="",
+                server_cert_path='',
                 notification_handler=mock_notification_handler,
                 reconnection_interval=60,
-            pod_id='localhost'
+                pod_id='localhost',
+            )
+        ]
     )
-    ])
 
     server1.connect.assert_called_once()
     server2.connect.assert_called_once()
 
-    mock_send_notification.assert_has_calls([
+    mock_send_notification.assert_has_calls(
+        [
             call(
                 metadata={
                     'model_id': '-',
                     'model_name': '-',
                     'workflow_name': '-',
-                'schedule_name': 'INITIALIZATION'
+                    'schedule_name': 'INITIALIZATION',
                 },
-            notification_id="OPC_CONNECTION_ERROR_server3",
+                notification_id='OPC_CONNECTION_ERROR_server3',
-            message="Failed to connect to OPC server: Test error",
+                message='Failed to connect to OPC server: Test error',
-            block="opc_repository",
+                block='opc_repository',
                 level=NotificationLevel.ERROR,
-            attachment_content=ANY
+                attachment_content=ANY,
+            )
+        ]
     )
-    ])
 
 
 @pytest_asyncio.fixture
-@patch("laborious.activities.opc.OpcRepository")
+@patch('laborious.activities.opc.OpcRepository')
 async def opc(mock_opc_repository):
     servers = {
         'server1': {
@@ -161,17 +163,9 @@ async def opc(mock_opc_repository):
         }
     }
 
-    mock_opc_repository.return_value.write_data = AsyncMock(
+    mock_opc_repository.return_value.write_data = AsyncMock(return_value=(True, {}))
-        return_value=(True, {})
+    mock_opc_repository.return_value.connect = AsyncMock(return_value=(True, {}))
-    )
+    opc = OPC(opc_servers=servers, logger=MagicMock(), notification_handler=MagicMock())
-    mock_opc_repository.return_value.connect = AsyncMock(
-        return_value=(True, {})
-    )
-    opc = OPC(
-        opc_servers=servers,
-        logger=MagicMock(),
-        notification_handler=MagicMock()
-    )
     await opc.init_opc()
     opc.send_notification = MagicMock()
     return opc
@@ -188,58 +182,79 @@ WRITE_DATA_CASES = [
 @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
 @mark.asyncio
 async def test_write_data_success(opc, tag, data_type, data):
-    result = await opc.write_data(server_id='server1', tag=tag, data=data,
+    result = await opc.write_data(
-                                  data_type=data_type, tag_type='prediction', metadata=metadata)
+        server_id='server1',
+        tag=tag,
+        data=data,
+        data_type=data_type,
+        tag_type='prediction',
+        metadata=metadata,
+    )
     assert result is True
     opc.opc_repository['server1'].write_data.assert_called_once_with(
-        tag, data, data_type, opc.logger, metadata)
+        tag, data, data_type, opc.logger, metadata
+    )
 
 
 @mark.asyncio
 async def test_write_data_failed(opc):
-    opc.opc_repository['server1'].write_data.return_value = (False, {
+    opc.opc_repository['server1'].write_data.return_value = (
+        False,
+        {
             'notification_id': 'OPC_WRITE_DATA_ERROR_server1',
             'message': 'Failed to write data to OPC server: Test error',
             'block': 'opc_repository',
             'level': NotificationLevel.ERROR,
-        'attachment_content': 'Test error'
+            'attachment_content': 'Test error',
-    })
+        },
+    )
 
-    result = await opc.write_data(server_id='server1', tag='tag1', data=50,
+    result = await opc.write_data(
-                                  data_type='int', tag_type='prediction', metadata=metadata)
+        server_id='server1',
+        tag='tag1',
+        data=50,
+        data_type='int',
+        tag_type='prediction',
+        metadata=metadata,
+    )
     assert result is False
 
     opc.send_notification.assert_called_once_with(
         metadata=metadata,
-        notification_id="OPC_WRITE_DATA_ERROR_server1",
+        notification_id='OPC_WRITE_DATA_ERROR_server1',
-        message="Failed to write data to OPC server: Test error",
+        message='Failed to write data to OPC server: Test error',
-        block="opc_repository",
+        block='opc_repository',
         level=NotificationLevel.ERROR,
-        attachment_content=ANY
+        attachment_content=ANY,
     )
 
 
 @mark.asyncio
 async def test_write_data_exception(opc):
-    opc.opc_repository['server1'].write_data.side_effect = Exception(
+    opc.opc_repository['server1'].write_data.side_effect = Exception('Test error')
-        "Test error")
 
     try:
-        await opc.write_data(server_id='server1', tag='tag1', data=50,
+        await opc.write_data(
-                             data_type='int', tag_type='prediction', metadata=metadata)
+            server_id='server1',
+            tag='tag1',
+            data=50,
+            data_type='int',
+            tag_type='prediction',
+            metadata=metadata,
+        )
 
     except Exception:
         opc.send_notification.assert_called_once_with(
             metadata=metadata,
-            notification_id="WRITE_OPC_PREDICTION_ERROR",
+            notification_id='WRITE_OPC_PREDICTION_ERROR',
-            message="Error writing data to OPC server: Test error",
+            message='Error writing data to OPC server: Test error',
-            block="write_opc_data",
+            block='write_opc_data',
             level=NotificationLevel.ERROR,
-            attachment_content=ANY
+            attachment_content=ANY,
         )
 
     else:
-        assert False, "Expected an exception to be raised"
+        raise AssertionError('Expected an exception to be raised')
 
 
 @mark.asyncio
@@ -247,20 +262,13 @@ async def test_write_opc_data_success(opc):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
         'opc_output_config': {
             'server1': {
-                'prediction_tags': {
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
-                    'tag1': {'data_type': 'float'}
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
+            }
         },
-                'confidence_tags': {
-                    'tag2': {'data_type': 'float'}
-                }
-            }
-        }
     }
 
     # Act
@@ -270,25 +278,30 @@ async def test_write_opc_data_success(opc):
 
     # Assert
     assert output == {'data': 'data'}
-    opc.write_data.assert_has_calls([
+    opc.write_data.assert_has_calls(
+        [
             call(
                 server_id='server1',
                 tag='tag1',
                 data=0.75,
                 data_type='float',
                 tag_type='prediction',
-            metadata=metadata['metadata']
+                metadata=metadata['metadata'],
-        )])
+            )
-    opc.write_data.assert_has_calls([
+        ]
+    )
+    opc.write_data.assert_has_calls(
+        [
             call(
                 server_id='server1',
                 tag='tag2',
                 data=0.95,
                 data_type='float',
                 tag_type='confidence',
-            metadata=metadata['metadata']
+                metadata=metadata['metadata'],
+            )
+        ]
     )
-    ])
     assert opc.write_data.call_count == 2
 
 
@@ -297,17 +310,9 @@ async def test_write_opc_data_empty_config(opc):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
         'opc_servers': ['server1'],
-        'opc_output_config': {
+        'opc_output_config': {'server1': {'prediction_tags': {}, 'confidence_tags': {}}},
-            'server1': {
-                'prediction_tags': {},
-                'confidence_tags': {}
-            }
-        }
     }
 
     # Act
@@ -322,20 +327,13 @@ async def test_write_opc_data_no_validate_server(opc):
     opc.validate_server = MagicMock(return_value=False)
     input_data = {
         **metadata,
-        'data': {
+        'data': {'prediction': [0.75], 'prediction_confidence': [0.95]},
-            'prediction': [0.75],
-            'prediction_confidence': [0.95]
-        },
         'opc_output_config': {
             'server1': {
-                'prediction_tags': {
+                'prediction_tags': {'tag1': {'data_type': 'float'}},
-                    'tag1': {'data_type': 'float'}
+                'confidence_tags': {'tag2': {'data_type': 'float'}},
+            }
         },
-                'confidence_tags': {
-                    'tag2': {'data_type': 'float'}
-                }
-            }
-        }
     }
 
     # Act
@@ -345,10 +343,13 @@ async def test_write_opc_data_no_validate_server(opc):
     opc.opc_repository['server1'].write_data.assert_not_called()
 
 
-@mark.parametrize('data,success,expected', [
+@mark.parametrize(
+    'data,success,expected',
+    [
         (DataFrame({'prediction_confidence': [0]}), True, 0),
         (DataFrame({'prediction_confidence': [0]}), False, 12),
-])
+    ],
+)
 def test_process_confidence(opc, data, success, expected):
     # Act
     result = opc.process_confidence(data, success, metadata)

tests/laborious/activities/test_storage.py

@@ -0,0 +1,213 @@
+import datetime
+from unittest.mock import ANY, AsyncMock, MagicMock, patch
+
+from pytest import fixture, mark, raises
+from sientia_do.notifications.models import NotificationLevel
+from sientia_do.temporal.activities.postgres import Postgres
+
+from laborious.activities.storage import Storage
+
+metadata = {
+    'metadata': {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'test_workflow',
+        'schedule_name': 'test_schedule',
+    }
+}
+
+
+@fixture
+@patch('laborious.activities.storage.MinioRepository')
+def storage(mock_minio_repository):
+    return Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        minio_config={
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___not_hasattr(mock_minio_repository):
+    logger = MagicMock()
+    notification_handler = MagicMock()
+    storage = Storage(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        minio_config={
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
+        logger=logger,
+        notification_handler=notification_handler,
+    )
+    assert isinstance(storage, Postgres)
+
+    mock_minio_repository.assert_called_once_with(
+        logger=logger,
+        notification_handler=notification_handler,
+        minio_endpoint_url='localhost:9000',
+        minio_access_key='minio',
+        minio_secret_key='minio123',
+        minio_region_name='us-east-1',
+        minio_default_bucket='test',
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___none_minio_repository(mock_minio_repository, storage):
+    storage.minio_repository = None
+    logger = MagicMock()
+    notification_handler = MagicMock()
+
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        minio_config={
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
+        logger=logger,
+        notification_handler=notification_handler,
+    )
+
+    mock_minio_repository.assert_called_once_with(
+        logger=logger,
+        notification_handler=notification_handler,
+        minio_endpoint_url='localhost:9000',
+        minio_access_key='minio',
+        minio_secret_key='minio123',
+        minio_region_name='us-east-1',
+        minio_default_bucket='test',
+    )
+
+
+@patch('laborious.activities.storage.MinioRepository')
+def test___init___done_repository(mock_minio_repository, storage):
+    storage.__init__(
+        host='localhost',
+        port=5432,
+        user='postgres',
+        password='postgres',
+        dbname='postgres',
+        min_connections=1,
+        max_connections=10,
+        minio_config={
+            'endpoint_url': 'localhost:9000',
+            'access_key': 'minio',
+            'secret_key': 'minio123',
+            'region_name': 'us-east-1',
+            'default_bucket': 'test',
+        },
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+    )
+    mock_minio_repository.assert_not_called()
+    assert storage.minio_repository is not None
+
+
+@mark.asyncio
+async def test_query_to_minio_minio_repository_not_initialized(storage):
+    storage.minio_repository = None
+
+    with raises(ValueError) as e:
+        await storage.query_to_minio({})
+
+        assert str(e.value) == 'Minio repository not initialized'
+
+
+@mark.asyncio
+async def test_query_to_minio_not_data(storage):
+    storage.load_custom_query = AsyncMock(return_value=None)
+    result = await storage.query_to_minio({})
+
+    storage.load_custom_query.assert_called_once_with({})
+    assert result['success'] is False
+    assert result['message'] == 'No data returned from query'
+
+
+@mark.asyncio
+@patch('laborious.activities.storage.pd.DataFrame')
+@patch('laborious.activities.storage.now')
+async def test_query_to_minio_success(now, dataframe, storage):
+    data = [{'a': 1}, {'a': 2}, {'a': 3}]
+    storage.load_custom_query = AsyncMock(return_value=data)
+    now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
+    storage.minio_repository.minio_bucket = 'test'
+
+    result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
+
+    dataframe.assert_called_once_with(data)
+
+    storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
+        dataframe=dataframe.return_value,
+        uri='s3://test/test_2024-01-01_00-00-00.parquet',
+        object_name='test_2024-01-01_00-00-00.parquet',
+        metadata=metadata['metadata'],
+    )
+
+    assert result['success'] is True
+    assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
+    assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'
+
+
+@mark.asyncio
+async def test_query_to_minio_error(storage):
+    storage.send_notification = MagicMock()
+    storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
+    result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
+    assert result['success'] is False
+    assert result['message'] == 'test'
+    storage.send_notification.assert_called_once_with(
+        metadata=metadata['metadata'],
+        notification_id='ERROR_STORING_QUERY_TO_MINIO',
+        message='Error storing query to MinIO: test',
+        block='query_to_minio',
+        level=NotificationLevel.ERROR,
+        attachment_content=ANY,
+    )
+
+
+def test_close(storage):
+    storage.minio_repository = MagicMock()
+
+    storage.close()
+
+    assert storage.minio_repository is None
+
+
+def test___del__(storage):
+    storage.close = MagicMock()
+
+    storage.__del__()
+
+    storage.close.assert_called_once()

tests/laborious/utils/filters/test_conditional_filters.py

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

tests/laborious/utils/filters/test_mlflow_filters.py

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

tests/laborious/utils/repository/test_minio_repository.py

@@ -0,0 +1,139 @@
+from unittest.mock import MagicMock, patch
+
+from botocore.utils import ClientError
+from pytest import fixture, raises
+
+from laborious.utils.repository.minio_repository import MinioRepository
+
+
+@patch('laborious.utils.repository.minio_repository.boto3')
+@patch('laborious.utils.repository.minio_repository.Config')
+def test___init___(mock_config, mock_boto3):
+    minio_repository = MinioRepository(
+        minio_endpoint_url='localhost:9000',
+        minio_access_key='minio',
+        minio_secret_key='minio123',
+        minio_region_name='us-east-1',
+        minio_default_bucket='test',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+    )
+
+    assert minio_repository.storage_options == {
+        'key': 'minio',
+        'secret': 'minio123',
+        'client_kwargs': {'endpoint_url': 'localhost:9000'},
+    }
+    assert minio_repository.minio_bucket == 'test'
+    assert minio_repository.minio_endpoint_url == 'localhost:9000'
+    assert minio_repository.minio_region_name == 'us-east-1'
+
+    mock_config.assert_called_once_with(
+        signature_version='s3v4',
+        s3={'addressing_style': 'path'},
+        retries={'max_attempts': 5, 'mode': 'standard'},
+        connect_timeout=5,
+        read_timeout=120,
+    )
+
+    mock_boto3.client.assert_called_once_with(
+        's3',
+        endpoint_url='localhost:9000',
+        aws_access_key_id='minio',
+        aws_secret_access_key='minio123',
+        region_name='us-east-1',
+        config=mock_config.return_value,
+    )
+
+
+@fixture
+@patch('laborious.utils.repository.minio_repository.Config')
+@patch('laborious.utils.repository.minio_repository.boto3')
+def minio_repository(mock_boto3, mock_config):
+    return MinioRepository(
+        minio_endpoint_url='localhost:9000',
+        minio_access_key='minio',
+        minio_secret_key='minio123',
+        minio_region_name='us-east-1',
+        minio_default_bucket='test',
+        logger=MagicMock(),
+        notification_handler=MagicMock(),
+    )
+
+
+def test_close(minio_repository):
+    minio_repository.close()
+    minio_repository.s3_client.close.assert_called_once()
+
+
+def test_ensure_bucket_exists_bucket_exists(minio_repository):
+    assert minio_repository.ensure_bucket_exists({}) is None
+
+    minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
+
+
+def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository):
+    minio_repository.s3_client.head_bucket.side_effect = ClientError(
+        error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
+    )
+
+    assert minio_repository.ensure_bucket_exists({}) is None
+
+    minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
+    minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
+
+
+def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository):
+    minio_repository.send_notification = MagicMock()
+
+    minio_repository.s3_client.head_bucket.side_effect = ClientError(
+        error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
+    )
+    minio_repository.s3_client.create_bucket.side_effect = ClientError(
+        error_response={'Error': {'Code': '404'}}, operation_name='create_bucket'
+    )
+
+    with raises(ClientError):
+        minio_repository.ensure_bucket_exists({})
+
+        minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
+        minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
+
+
+@patch('laborious.utils.repository.minio_repository.BytesIO')
+def test_store_dataframe_as_parquet(mock_bytesio, minio_repository):
+    input_data = MagicMock()
+
+    minio_repository.ensure_bucket_exists = MagicMock(return_value=True)
+
+    minio_repository.store_dataframe_as_parquet(
+        dataframe=input_data, uri='s3://test/test.parquet', object_name='test.parquet', metadata={}
+    )
+
+    minio_repository.ensure_bucket_exists.assert_called_once_with({})
+    mock_bytesio.assert_called_once()
+
+    input_data.to_parquet.assert_called_once_with(
+        mock_bytesio.return_value, engine='pyarrow', index=True
+    )
+    mock_bytesio.return_value.seek.assert_called_once_with(0)
+    minio_repository.s3_client.put_object.assert_called_once_with(
+        Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value
+    )
+
+
+@patch('laborious.utils.repository.minio_repository.BytesIO')
+@patch('laborious.utils.repository.minio_repository.read_parquet')
+def test_get_parquet_as_dataframe(mock_read_parquet, mock_bytesio, minio_repository):
+    input_data = {'Body': MagicMock(read=MagicMock(return_value=b'test'))}
+
+    minio_repository.s3_client.get_object.return_value = input_data
+
+    output = minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={})
+
+    minio_repository.s3_client.get_object.assert_called_once_with(Bucket='test', Key='test.parquet')
+
+    mock_bytesio.assert_called_once_with(input_data['Body'].read.return_value)
+    mock_read_parquet.assert_called_once_with(mock_bytesio.return_value)
+
+    assert output == mock_read_parquet.return_value

tests/laborious/utils/repository/test_opc_repository.py

@@ -1,9 +1,11 @@
-import pytest
-from unittest.mock import AsyncMock, Mock, patch, MagicMock, ANY, call
-from asyncua.crypto.security_policies import SecurityPolicyBasic256
-from laborious.utils.repository.opc_repository import OpcRepository
-from sientia_do.notifications.models import NotificationLevel
 from datetime import datetime
+from unittest.mock import ANY, AsyncMock, MagicMock, Mock, call, patch
+
+import pytest
+from asyncua.crypto.security_policies import SecurityPolicyBasic256
+from sientia_do.notifications.models import NotificationLevel
+
+from laborious.utils.repository.opc_repository import OpcRepository
 
 
 @pytest.fixture
@@ -14,15 +16,15 @@ def mock_logger():
 @pytest.fixture
 def opc_repository(mock_logger):
     return OpcRepository(
-        id="test_repo",
+        opc_id='test_repo',
-        url="opc.tcp://localhost:4840",
+        url='opc.tcp://localhost:4840',
         logger=mock_logger,
         notification_handler=Mock(),
         reconnection_interval=60,
-        server_uri="urn:test:server",
+        server_uri='urn:test:server',
-        cert_path="/path/to/cert.pem",
+        cert_path='/path/to/cert.pem',
-        private_key_path="/path/to/key.pem",
+        private_key_path='/path/to/key.pem',
-        server_cert_path="/path/to/server_cert.pem"
+        server_cert_path='/path/to/server_cert.pem',
     )
 
 
@@ -35,22 +37,22 @@ def mock_client():
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
 
 def test_init(opc_repository):
-    assert opc_repository.id == "test_repo"
+    assert opc_repository.id == 'test_repo'
-    assert opc_repository.url == "opc.tcp://localhost:4840"
+    assert opc_repository.url == 'opc.tcp://localhost:4840'
-    assert opc_repository.server_uri == "urn:test:server"
+    assert opc_repository.server_uri == 'urn:test:server'
-    assert opc_repository.cert_path == "/path/to/cert.pem"
+    assert opc_repository.cert_path == '/path/to/cert.pem'
-    assert opc_repository.private_key_path == "/path/to/key.pem"
+    assert opc_repository.private_key_path == '/path/to/key.pem'
-    assert opc_repository.server_cert_path == "/path/to/server_cert.pem"
+    assert opc_repository.server_cert_path == '/path/to/server_cert.pem'
     assert opc_repository.reconnection_interval == 60
     assert opc_repository.client is None
     assert opc_repository.last_reconnection_time is None
@@ -62,12 +64,12 @@ async def test_set_security(opc_repository, mock_client):
     opc_repository.client = mock_client
     await opc_repository.set_security()
 
-    mock_client.application_uri = "urn:test:server"
+    mock_client.application_uri = 'urn:test:server'
     mock_client.set_security.assert_called_once_with(
         SecurityPolicyBasic256,
-        certificate="/path/to/cert.pem",
+        certificate='/path/to/cert.pem',
-        private_key="/path/to/key.pem",
+        private_key='/path/to/key.pem',
-        server_certificate="/path/to/server_cert.pem"
+        server_certificate='/path/to/server_cert.pem',
     )
     assert mock_client.secure_channel_timeout == 10000000
     assert mock_client.session_timeout == 10000000
@@ -81,8 +83,16 @@ async def test_set_security_missing_certificates(opc_repository):
     try:
         await opc_repository.set_security()
     except ValueError as e:
-        assert str(
+        assert str(e) == 'Certificate and private key paths must be provided for secure connection.'
-            e) == "Certificate and private key paths must be provided for secure connection."
+
+
+@pytest.mark.asyncio
+async def test_set_security_missing_client(opc_repository):
+    opc_repository.client = None
+    try:
+        await opc_repository.set_security()
+    except ValueError as e:
+        assert str(e) == 'Client must be initialized before setting security'
 
 
 @pytest.mark.asyncio
@@ -123,19 +133,34 @@ async def test_try_connect_success(opc_repository):
 async def test_try_connect_fail(opc_repository):
     opc_repository.last_reconnection_time = None
     opc_repository.client = MagicMock()
-    opc_repository.client.connect.side_effect = Exception("Test error")
+    opc_repository.client.connect.side_effect = Exception('Test error')
 
     is_connected, error_data = await opc_repository.try_connect()
 
     opc_repository.client.connect.assert_called_once()
     assert is_connected is False
-    assert error_data['notification_id'] == f"OPC_CONNECTION_ERROR_{opc_repository.id}"
+    assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'
-    assert error_data['message'] == "Failed to connect to OPC server: Test error"
+    assert error_data['message'] == 'Failed to connect to OPC server: Test error'
-    assert error_data['block'] == "opc_repository"
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None
 
 
+@pytest.mark.asyncio
+async def test_try_connect_no_client(opc_repository):
+    opc_repository.client = None
+    result = await opc_repository.try_connect()
+    assert result == (
+        False,
+        {
+            'notification_id': f'OPC_CONNECTION_ERROR_{opc_repository.id}',
+            'message': 'Client is not initialized',
+            'block': 'opc_repository',
+            'level': NotificationLevel.ERROR,
+        },
+    )
+
+
 @pytest.mark.asyncio
 async def test_disconnect(opc_repository, mock_client):
     opc_repository.client = mock_client
@@ -154,12 +179,11 @@ async def test_disconnect_no_client(opc_repository):
 @pytest.mark.asyncio
 async def test_disconnect_error(opc_repository, mock_client):
     opc_repository.client = mock_client
-    mock_client.disconnect.side_effect = Exception("Test error")
+    mock_client.disconnect.side_effect = Exception('Test error')
     await opc_repository.disconnect()
 
     opc_repository.logger.custom_error.assert_called_once_with(
-        "Failed to disconnect from OPC server: Test error",
+        'Failed to disconnect from OPC server: Test error', ANY
-        ANY
     )
     assert opc_repository.client is None
 
@@ -177,9 +201,7 @@ async def test_validate_connection_none_client(opc_repository):
 async def test_validate_connection_error_count_disconnect_error(opc_repository):
     opc_repository.error_count = 6
     opc_repository.client = AsyncMock()
-    opc_repository.disconnect = AsyncMock(
+    opc_repository.disconnect = AsyncMock(side_effect=Exception('Test error'))
-        side_effect=Exception("Test error")
-    )
     opc_repository.connect = AsyncMock(return_value=(True, {}))
 
     response = await opc_repository.validate_connection()
@@ -188,34 +210,34 @@ async def test_validate_connection_error_count_disconnect_error(opc_repository):
     opc_repository.connect.assert_called_once()
     opc_repository.logger.custom_error.assert_has_calls(
         [
-            call("Failed to disconnect from OPC server: Test error", ANY),
+            call('Failed to disconnect from OPC server: Test error', ANY),
         ]
     )
 
 
 @pytest.mark.asyncio
 async def test_validate_connection_error_validate_connection_error(opc_repository):
-    opc_repository.client = MagicMock(
+    opc_repository.client = MagicMock(uaclient=Exception('Test error'))
-        uaclient=Exception("Test error")
-    )
     opc_repository.error_count = 0
 
     response = await opc_repository.validate_connection()
 
-    assert response == (False, {
+    assert response == (
-        "notification_id": f"OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}",
+        False,
-        "message": "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
+        {
-        "block": "opc_repository",
+            'notification_id': f'OPC_CONNECTION_CHECK_ERROR_{opc_repository.id}',
-        "level": NotificationLevel.ERROR,
+            'message': "Failed to validate connection to OPC server: 'Exception' object has no attribute 'protocol'",
-        "attachment_content": ANY
+            'block': 'opc_repository',
-    })
+            'level': NotificationLevel.ERROR,
+            'attachment_content': ANY,
+        },
+    )
 
 
 @pytest.mark.asyncio
 @patch('laborious.utils.repository.opc_repository.datetime')
 async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, opc_repository):
-    _mock_datetime.now = MagicMock(
+    _mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 0, 0, 0))
-        return_value=datetime(2025, 1, 1, 0, 0, 0))
     opc_repository.error_count = 0
     opc_repository.client = MagicMock()
     opc_repository.client.uaclient.protocol = None
@@ -224,19 +246,21 @@ async def test_validate_connection_lost_not_time_to_reconnect(_mock_datetime, op
 
     response = await opc_repository.validate_connection()
     opc_repository.connect.assert_not_called()
-    assert response == (False, {
+    assert response == (
-        "notification_id": f"OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}",
+        False,
-        "message": f"OPC server {opc_repository.id} is not connected, waiting for next reconnection window...",
+        {
-        "block": "opc_repository",
+            'notification_id': f'OPC_CONNECTION_AWAITING_RECONNECTION_WINDOW_{opc_repository.id}',
-        "level": NotificationLevel.WARNING
+            'message': f'OPC server {opc_repository.id} is not connected, waiting for next reconnection window...',
-    })
+            'block': 'opc_repository',
+            'level': NotificationLevel.WARNING,
+        },
+    )
 
 
 @pytest.mark.asyncio
 @patch('laborious.utils.repository.opc_repository.datetime')
 async def test_validate_connection_lost_time_to_reconnect(mock_datetime, opc_repository):
-    mock_datetime.now = MagicMock(
+    mock_datetime.now = MagicMock(return_value=datetime(2025, 1, 1, 1, 0, 0))
-        return_value=datetime(2025, 1, 1, 1, 0, 0))
     opc_repository.error_count = 0
     opc_repository.client = AsyncMock()
     opc_repository.client.uaclient.protocol = None
@@ -253,7 +277,7 @@ async def test_validate_connection_success(opc_repository):
     opc_repository.client = MagicMock()
     opc_repository.error_count = 0
     opc_repository.client.uaclient.protocol = MagicMock()
-    opc_repository.client.uaclient.protocol.state = "open"
+    opc_repository.client.uaclient.protocol.state = 'open'
 
     output = await opc_repository.validate_connection()
     assert output == (True, {})
@@ -262,17 +286,16 @@ async def test_validate_connection_success(opc_repository):
 @pytest.mark.asyncio
 async def test_write_data_validate_connection_do_nothing(opc_repository):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
-    opc_repository.client = AsyncMock(
+    opc_repository.client = AsyncMock(get_node=MagicMock())
-        get_node=MagicMock()
-    )
     mock_node = AsyncMock()
     opc_repository.client.get_node.return_value = mock_node
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    result = await opc_repository.write_data(
-                                             "float", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
     assert result == (True, {})
 
 
@@ -282,8 +305,9 @@ async def test_write_data_validate_connection_failed(opc_repository):
     opc_repository.client = AsyncMock()
     opc_repository.error_count = 0
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    result = await opc_repository.write_data(
-                                             "float", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_not_called()
@@ -295,18 +319,21 @@ async def test_write_data_get_node_failed(opc_repository):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
     opc_repository.client = AsyncMock()
     opc_repository.error_count = 0
-    opc_repository.client.get_node = MagicMock(
+    opc_repository.client.get_node = MagicMock(side_effect=Exception('Test error'))
-        side_effect=Exception("Test error"))
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    is_success, error_data = await opc_repository.write_data(
-                                                             "float", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    opc_repository.client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}"
+    assert error_data['notification_id'] == f'OPC_WRITE_GET_NODE_ERROR_{opc_repository.id}'
-    assert error_data['message'] == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    assert (
-    assert error_data['block'] == "opc_repository"
+        error_data['message']
+        == "Failed to get node from OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None
 
@@ -318,16 +345,20 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client):
     mock_node = AsyncMock()
     mock_client.get_node = MagicMock(return_value=mock_node)
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    is_success, error_data = await opc_repository.write_data(
-                                                             "invalid_type", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'invalid_type', opc_repository.logger, metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
 
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}"
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_TYPE_ERROR_{opc_repository.id}'
-    assert error_data['message'] == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    assert (
-    assert error_data['block'] == "opc_repository"
+        error_data['message']
+        == "Unsupported data type: invalid_type | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data.get('attachment_content') is None
 
@@ -340,10 +371,11 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
     mock_node = AsyncMock()
     mock_client.get_node = MagicMock(return_value=mock_node)
 
-    result = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    result = await opc_repository.write_data(
-                                             "float", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
+    )
 
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
     mock_node.write_value.assert_called_once()
     assert result == (True, {})
 
@@ -351,7 +383,7 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
         pod_id=opc_repository.pod_id,
         model_name=metadata['metadata']['model_name'],
         pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id
+        opc_server_id=opc_repository.id,
     )
     mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
 
@@ -359,10 +391,11 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
         pod_id=opc_repository.pod_id,
         model_name=metadata['metadata']['model_name'],
         pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id
+        opc_server_id=opc_repository.id,
     )
     mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
-        ANY)
+        ANY
+    )
 
 
 @pytest.mark.asyncio
@@ -372,17 +405,21 @@ async def test_write_data_write_value_failed(opc_repository, mock_client):
     mock_node = AsyncMock()
     opc_repository.error_count = 0
     mock_client.get_node = MagicMock(return_value=mock_node)
-    mock_node.write_value.side_effect = Exception("Test error")
+    mock_node.write_value.side_effect = Exception('Test error')
 
-    is_success, error_data = await opc_repository.write_data("ns=2;s=TestNode", 42.0,
+    is_success, error_data = await opc_repository.write_data(
-                                                             "float", opc_repository.logger, metadata['metadata'])
+        'ns=2;s=TestNode', 42.0, 'float', opc_repository.logger, metadata['metadata']
+    )
 
     opc_repository.validate_connection.assert_called_once()
-    mock_client.get_node.assert_called_once_with("ns=2;s=TestNode")
+    mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
     mock_node.write_value.assert_called_once()
     assert is_success is False
-    assert error_data['notification_id'] == f"OPC_WRITE_DATA_ERROR_{opc_repository.id}"
+    assert error_data['notification_id'] == f'OPC_WRITE_DATA_ERROR_{opc_repository.id}'
-    assert error_data['message'] == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    assert (
-    assert error_data['block'] == "opc_repository"
+        error_data['message']
+        == "Failed to write data to OPC server: Test error | metadata: {'model_id': 'test_model', 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule'}"
+    )
+    assert error_data['block'] == 'opc_repository'
     assert error_data['level'] == NotificationLevel.ERROR
     assert error_data['attachment_content'] is not None

tests/laborious/utils/test_connectors_config.py

@@ -1,8 +1,12 @@
 from os import environ
-from laborious.utils.connectors_config import (build_mlflow_config,
+
+from laborious.utils.connectors_config import (
+    build_minio_config,
+    build_mlflow_config,
+    build_mongodb_config,
     build_opc_config,
     build_postgres_config,
-                                               build_mongodb_config)
+)
 
 
 def test_build_mlflow_config_with_env_vars():
@@ -144,7 +148,7 @@ def test_build_mongo_db_config_with_env_vars():
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
         'database_name': 'test_db',
-        'ttl_index_seconds': 3600
+        'ttl_index_seconds': 3600,
     }
 
 
@@ -157,5 +161,35 @@ def test_build_mongo_db_config_with_defaults():
     assert build_mongodb_config() == {
         'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
         'database_name': 'sientia',
-        'ttl_index_seconds': 3600
+        'ttl_index_seconds': 3600,
+    }
+
+
+def test_build_minio_config_with_env_vars():
+    environ['MINIO_ENDPOINT_URL'] = 'http://test-host'
+    environ['MINIO_ACCESS_KEY'] = 'test-key'
+    environ['MINIO_SECRET_KEY'] = 'test-secret'
+    environ['MINIO_REGION_NAME'] = 'test-region'
+    environ['MINIO_DEFAULT_BUCKET'] = 'test-bucket'
+    assert build_minio_config() == {
+        'endpoint_url': 'http://test-host',
+        'access_key': 'test-key',
+        'secret_key': 'test-secret',
+        'region_name': 'test-region',
+        'default_bucket': 'test-bucket',
+    }
+
+
+def test_build_minio_config_with_defaults():
+    environ.pop('MINIO_ENDPOINT_URL', None)
+    environ.pop('MINIO_ACCESS_KEY', None)
+    environ.pop('MINIO_SECRET_KEY', None)
+    environ.pop('MINIO_REGION_NAME', None)
+    environ.pop('MINIO_DEFAULT_BUCKET', None)
+    assert build_minio_config() == {
+        'endpoint_url': 'http://localhost:9000',
+        'access_key': 'minioadmin',
+        'secret_key': 'minioadmin',
+        'region_name': 'us-east-1',
+        'default_bucket': 'laborious',
     }

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -1,9 +1,10 @@
-from unittest.mock import call, patch, AsyncMock, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
-from pytest import mark, fixture
+
+from pytest import fixture, mark
+from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from laborious.activities.activities import Activities
 from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction
-from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 
 @fixture
@@ -12,36 +13,39 @@ def format_and_export_prediction():
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
 async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
-
     input_data = {
         'metadata': metadata,
-        "path_flag": None,
+        'path_flag': None,
-        "data": {"test": "data"},
+        'data': {'test': 'data'},
-        "timestamp": "2021-01-01",
+        'timestamp': '2021-01-01',
-        "model_id": 1,
+        'model_id': 1,
-        "prediction_confidence": 0,
+        'prediction_confidence': 0,
-        "schema": "test_schema",
+        'schema': 'test_schema',
-        "table_name": "test_table",
+        'table_name': 'test_table',
-        "opc_servers": ["test_server"],
+        'opc_servers': ['test_server'],
-        "opc_output_config": {"test": "config"},
+        'opc_output_config': {'test': 'config'},
-        "prediction_store_policy": "erl:1"
+        'prediction_store_policy': 'erl:1',
     }
 
     await format_and_export_prediction.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
             call(
                 Activities.format_prediction,
                 {
@@ -50,26 +54,31 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                     'model_id': input_data['model_id'],
                     'prediction_confidence': input_data['prediction_confidence'],
                     'prediction_store_policy': input_data['prediction_store_policy'],
-                **metadata
+                    **metadata,
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
-        )])
+            )
+        ]
+    )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.write_opc_data,
                 {
                     'opc_output_config': input_data['opc_output_config'],
                     'data': workflow_mock.execute_local_activity_method.return_value,
-                **metadata
+                    **metadata,
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.export_data_to_postgres,
                 {
@@ -79,38 +88,43 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
+                        'format': DATETIME_FORMAT_WITH_TZ,
-                }
+                    },
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
-        )])
+            )
+        ]
+    )
 
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.format_and_export_prediction.workflow", new_callable=AsyncMock)
+@patch(
+    'laborious.workflows.sub_workflows.format_and_export_prediction.workflow',
+    new_callable=AsyncMock,
+)
 async def test_run_default_path_flag(workflow_mock, format_and_export_prediction):
-
     input_data = {
         'metadata': metadata,
-        "path_flag": "default",
+        'path_flag': 'default',
-        "data": {"test": "data"},
+        'data': {'test': 'data'},
-        "timestamp": "2021-01-01",
+        'timestamp': '2021-01-01',
-        "model_id": 1,
+        'model_id': 1,
-        "prediction_confidence": 0,
+        'prediction_confidence': 0,
-        "schema": "test_schema",
+        'schema': 'test_schema',
-        "table_name": "test_table",
+        'table_name': 'test_table',
-        "opc_servers": ["test_server"],
+        'opc_servers': ['test_server'],
-        "opc_output_config": {"test": "config"},
+        'opc_output_config': {'test': 'config'},
-        "comment": "test_comment"
+        'comment': 'test_comment',
     }
 
     await format_and_export_prediction.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
             call(
                 Activities.format_default_prediction,
                 {
@@ -118,27 +132,31 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                     'model_id': input_data['model_id'],
                     'prediction_confidence': input_data['prediction_confidence'],
                     'comment': input_data['comment'],
-                **metadata
+                    **metadata,
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.write_opc_data,
                 {
                     'opc_output_config': input_data['opc_output_config'],
                     'data': workflow_mock.execute_local_activity_method.return_value,
-                **metadata
+                    **metadata,
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.export_data_to_postgres,
                 {
@@ -148,13 +166,14 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
-                    'format': DATETIME_FORMAT_WITH_TZ
+                        'format': DATETIME_FORMAT_WITH_TZ,
-                }
+                    },
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, patch, call, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from laborious.activities.activities import Activities
 from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
 
@@ -10,17 +12,17 @@ def prediction_process():
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model",
+        'model_id': 'test_model',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "test_workflow",
+        'workflow_name': 'test_workflow',
-        "schema_name": "test_schedule",
+        'schema_name': 'test_schedule',
     },
 }
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock, prediction_process):
     prediction_process.path_flag_handler = AsyncMock(return_value=False)
     # Arrange
@@ -34,26 +36,24 @@ async def test_run(workflow_mock, prediction_process):
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
+        'model_config': {'retention': '30'},
-            'retention': '30'
-        },
         'path_priority': ['continue', 'repeat', 'stop'],
         'opc_output_config': {'test': 'config'},
-        'prediction_store_policy': 'lts:1'
+        'prediction_store_policy': 'lts:1',
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
         {'content': 'predicted_data', 'timestamp': '2024-01-01'},  # request_predict
         # mlflow_response_gate (predict)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
     ]
 
     # Act
@@ -62,57 +62,112 @@ async def test_run(workflow_mock, prediction_process):
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 7
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
-        call(Activities.get_last_timestamp, {
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
                     **metadata,
                     'data': input_data['data'],
                 },
-            retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.input_gate, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
                     **metadata,
                     'filters': input_data['input_filters'],
                     'data': input_data['data'],
                     'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.request_transform, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
                     **metadata,
                     'data': input_data['data'],
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_response_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     **metadata,
                     'filters': input_data['mlflow_transform_filters'],
                     'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_content_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
                     **metadata,
                     'filters': input_data['mlflow_transform_filters'],
                     'data': 'transformed_data',
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.request_predict, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
                     **metadata,
                     'data': 'transformed_data',
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_response_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     **metadata,
                     'filters': input_data['mlflow_predict_filters'],
                     'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
                     'type': 'predict',
                     'path_priority': input_data['path_priority'],
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
 
     workflow_mock.execute_child_workflow.assert_called_once_with(
         'format_and_export_prediction',
@@ -129,13 +184,13 @@ async def test_run(workflow_mock, prediction_process):
             'schema': input_data['schema'],
             'table_name': input_data['table_name'],
             'comment': 'Error',
-            'prediction_store_policy': input_data['prediction_store_policy']
+            'prediction_store_policy': input_data['prediction_store_policy'],
-        }
+        },
     )
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
     prediction_process.path_flag_handler = AsyncMock(return_value=True)
     # Arrange
@@ -149,17 +204,15 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
+        'model_config': {'retention': '30'},
-            'retention': '30'
-        },
         'path_priority': ['continue', 'repeat', 'stop'],
-        'opc_output_config': {'test': 'config'}
+        'opc_output_config': {'test': 'config'},
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('stop', 0.95, "Input data with bad quality"),  # input_gate
+        ('stop', 0.95, 'Input data with bad quality'),  # input_gate
     ]
 
     # Act
@@ -167,23 +220,35 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process):
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 2
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
-        call(Activities.get_last_timestamp, {
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
                     'data': input_data['data'],
                     **metadata,
-             }, retry_policy=ANY, start_to_close_timeout=ANY),
+                },
-        call(Activities.input_gate, {
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+            call(
+                Activities.input_gate,
+                {
                     'filters': input_data['input_filters'],
                     'data': input_data['data'],
                     'path_priority': input_data['path_priority'],
                     **metadata,
-             }, retry_policy=ANY, start_to_close_timeout=ANY)
+                },
-    ])
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            ),
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process):
     prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True])
     # Arrange
@@ -197,19 +262,17 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
+        'model_config': {'retention': '30'},
-            'retention': '30'
-        },
         'path_priority': ['continue', 'repeat', 'stop'],
-        'opc_output_config': {'test': 'config'}
+        'opc_output_config': {'test': 'config'},
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('repeat', 0.95, "Input data with bad quality"),  # input_gate
+        ('repeat', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (transform)
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (transform)
     ]
 
     # Act
@@ -217,46 +280,72 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 4
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
-        call(Activities.get_last_timestamp, {
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
                     'data': input_data['data'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.input_gate, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
                     'filters': input_data['input_filters'],
                     'data': input_data['data'],
                     'path_priority': input_data['path_priority'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.request_transform, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
                     'data': input_data['data'],
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-            **metadata
+                    **metadata,
                 },
-            retry_policy=ANY, start_to_close_timeout=ANY)
+                retry_policy=ANY,
-    ])
+                start_to_close_timeout=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+            )
-        call(Activities.mlflow_response_gate, {
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     'filters': input_data['mlflow_transform_filters'],
                     'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
-            **metadata
+                    **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)
+                },
-    ])
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process):
-    prediction_process.path_flag_handler = AsyncMock(
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True])
-        side_effect=[False, False, True])
     # Arrange
     input_data = {
         'metadata': metadata,
@@ -268,22 +357,20 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
+        'model_config': {'retention': '30'},
-            'retention': '30'
-        },
         'path_priority': ['continue', 'repeat', 'stop'],
-        'opc_output_config': {'test': 'config'}
+        'opc_output_config': {'test': 'config'},
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
     ]
 
     # Act
@@ -292,51 +379,88 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 5
 
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
-        call(Activities.get_last_timestamp, {
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
                     'data': input_data['data'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.input_gate, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
                     'filters': input_data['input_filters'],
                     'data': input_data['data'],
                     'path_priority': input_data['path_priority'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.request_transform, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
                     'data': input_data['data'],
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-            **metadata
+                    **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_response_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     'filters': input_data['mlflow_transform_filters'],
                     'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
                     **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_content_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
                     'filters': input_data['mlflow_transform_filters'],
                     'data': 'transformed_data',
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
                     **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process):
-    prediction_process.path_flag_handler = AsyncMock(
+    prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True])
-        side_effect=[False, False, False, True])
     # Arrange
     input_data = {
         'metadata': metadata,
@@ -348,24 +472,22 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
         'mlflow_transform_filters': {'test': 'filter'},
         'mlflow_predict_filters': {'test': 'filter'},
         'model_name': 'test_model_name',
-        'model_config': {
+        'model_config': {'retention': '30'},
-            'retention': '30'
-        },
         'path_priority': ['continue', 'repeat', 'stop'],
-        'opc_output_config': {'test': 'config'}
+        'opc_output_config': {'test': 'config'},
     }
 
     # Mock the activity responses
     workflow_mock.execute_local_activity_method.side_effect = [
         '2024-01-01',  # get_last_timestamp
-        ('continue', 0.95, "Input data with bad quality"),  # input_gate
+        ('continue', 0.95, 'Input data with bad quality'),  # input_gate
         {'content': 'transformed_data', 'timestamp': '2024-01-01'},  # transform_data
         # mlflow_response_gate (transform)
-        ('continue', 0.95, "Error"),
+        ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
-        ('continue', 0.95, "Transformed data not passed the content filter"),
+        ('continue', 0.95, 'Transformed data not passed the content filter'),
         {'content': 'predicted_data', 'timestamp': '2024-01-01'},  # request_predict
-        ('continue', 0.95, "Error"),  # mlflow_response_gate (predict)
+        ('continue', 0.95, 'Error'),  # mlflow_response_gate (predict)
     ]
 
     # Act
@@ -373,63 +495,117 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
 
     # Assert
     assert workflow_mock.execute_local_activity_method.call_count == 7
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
-        call(Activities.get_last_timestamp, {
+        [
+            call(
+                Activities.get_last_timestamp,
+                {
                     'data': input_data['data'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.input_gate, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.input_gate,
+                {
                     'filters': input_data['input_filters'],
                     'data': input_data['data'],
                     'path_priority': input_data['path_priority'],
                     **metadata,
                 },
-             retry_policy=ANY, start_to_close_timeout=ANY)])
+                retry_policy=ANY,
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                start_to_close_timeout=ANY,
-        call(Activities.request_transform, {
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_transform,
+                {
                     'data': input_data['data'],
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-            **metadata
+                    **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_response_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     'filters': input_data['mlflow_transform_filters'],
                     'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
                     **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_content_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_content_gate,
+                {
                     'filters': input_data['mlflow_transform_filters'],
                     'data': 'transformed_data',
                     'type': 'transform',
                     'path_priority': input_data['path_priority'],
                     **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.request_predict, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.request_predict,
+                {
                     'data': 'transformed_data',
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
-            **metadata
+                    **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
-    workflow_mock.execute_local_activity_method.assert_has_calls([
+                retry_policy=ANY,
-        call(Activities.mlflow_response_gate, {
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.mlflow_response_gate,
+                {
                     'filters': input_data['mlflow_predict_filters'],
                     'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
                     'type': 'predict',
                     'path_priority': input_data['path_priority'],
                     **metadata,
-        }, retry_policy=ANY, start_to_close_timeout=ANY)])
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_stop(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -440,21 +616,24 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
+    model_config = {'retention': '30'}
-        'retention': '30'
-    }
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
-            'model_config': model_config
+            'model_config': model_config,
-        }, confidence, last_timestamp, ""
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert
@@ -464,7 +643,7 @@ async def test_path_flag_handler_stop(workflow_mock, prediction_process):
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -475,21 +654,24 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
+    model_config = {'retention': '30'}
-        'retention': '30'
-    }
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
             'model_id': model,
             'last_timestamp': last_timestamp,
             'model_name': model_name,
-            'model_config': model_config
+            'model_config': model_config,
-        }, confidence, last_timestamp, ""
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert
@@ -504,13 +686,13 @@ async def test_path_flag_handler_repeat(workflow_mock, prediction_process):
             'last_timestamp': last_timestamp,
         },
         retry_policy=ANY,
-        start_to_close_timeout=ANY
+        start_to_close_timeout=ANY,
     )
     workflow_mock.execute_child_workflow.assert_not_called()
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -521,14 +703,14 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
+    model_config = {'retention': '30'}
-        'retention': '30'
-    }
     prediction_store_policy = 'erl:1'
 
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             'metadata': metadata,
             'schema': schema,
             'table_name': table_name,
@@ -537,8 +719,11 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'model_name': model_name,
             'model_config': model_config,
             'opc_output_config': {'test': 'config'},
-            'prediction_store_policy': prediction_store_policy
+            'prediction_store_policy': prediction_store_policy,
-        }, confidence, last_timestamp, 'Prediction Process'
+        },
+        confidence,
+        last_timestamp,
+        'Prediction Process',
     )
 
     # Assert
@@ -559,13 +744,13 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process):
             'table_name': table_name,
             'comment': 'Prediction Process',
             'opc_output_config': {'test': 'config'},
-            'prediction_store_policy': prediction_store_policy
+            'prediction_store_policy': prediction_store_policy,
-        }
+        },
     )
 
 
 @mark.asyncio
-@patch("laborious.workflows.sub_workflows.prediction_process.workflow", new_callable=AsyncMock)
+@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock)
 async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
     # Arrange
     data = {'test': 'data'}
@@ -576,13 +761,13 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
     model = 'test_model'
     last_timestamp = '2024-01-01'
     model_name = 'test_model_name'
-    model_config = {
+    model_config = {'retention': '30'}
-        'retention': '30'
-    }
     prediction_store_policy = 'erl:1'
     # Act
     result = await prediction_process.path_flag_handler(
-        data, path_flag, {
+        data,
+        path_flag,
+        {
             **metadata,
             'schema': schema,
             'table_name': table_name,
@@ -591,8 +776,11 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process):
             'model_name': model_name,
             'model_config': model_config,
             'opc_output_config': {'test': 'config'},
-            'prediction_store_policy': prediction_store_policy
+            'prediction_store_policy': prediction_store_policy,
-        }, confidence, last_timestamp, ""
+        },
+        confidence,
+        last_timestamp,
+        '',
     )
 
     # Assert

tests/laborious/workflows/test_minimal_retrain.py

@@ -1,5 +1,7 @@
-from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
+from unittest.mock import ANY, AsyncMock, call, patch
+
 from pytest import fixture, mark
+
 from laborious.activities.activities import Activities
 from laborious.workflows.minimal_retrain import MinimalRetrain
 
@@ -10,11 +12,11 @@ def minimal_retrain() -> MinimalRetrain:
 
 
 metadata = {
-    "metadata": {
+    'metadata': {
-        "model_id": "test_model_id",
+        'model_id': 'test_model_id',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "minimal_retrain",
+        'workflow_name': 'minimal_retrain',
-        "schedule_name": "test_schedule",
+        'schedule_name': 'test_schedule',
     },
 }
 
@@ -23,76 +25,273 @@ metadata = {
 @patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
 async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
     input_data = {
-        "model_id": "test_model_id",
+        'model_id': 'test_model_id',
-        "model_name": "test_model",
+        'model_name': 'test_model',
-        "workflow_name": "minimal_retrain",
+        'workflow_name': 'minimal_retrain',
-        "schedule_name": "test_schedule",
+        'schedule_name': 'test_schedule',
-        "query": "test_query",
+        'query': 'test_query',
-        "schema": "test_schema",
+        'schema': 'test_schema',
-        "table_name": "test_table",
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
     }
 
     workflow_mock.execute_activity_method = AsyncMock(
-        return_value={
+        side_effect=[
-            "data1": "1",
+            {'success': True, 'object_key': 'test_object_key'},
-            "data2": "2",
+            {'success': True, 'experiment': 'test_experiment'},
-        }
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
     )
 
     await minimal_retrain.run(input_data)
 
-    workflow_mock.execute_local_activity_method.assert_has_calls(
+    workflow_mock.execute_activity_method.assert_has_calls(
         [
             call(
-                Activities.load_custom_query,
+                Activities.query_to_minio,
                 {
                     **metadata,
-                    "query": input_data["query"],
+                    'query': input_data['query'],
-                    'datetime_columns': input_data.get('datetime_columns', [])
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                    'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
                 },
                 retry_policy=ANY,
-                start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
             )
         ]
     )
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.retrain_model,
                 {
                     **metadata,
-                'data': workflow_mock.execute_local_activity_method.return_value,
+                    'object_key': 'test_object_key',
                     'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.update_production_model,
                 {
                     **metadata,
                     'model_name': input_data['model_name'],
-                'model_id': input_data['model_id'],
+                    'success': True,
-                **workflow_mock.execute_activity_method.return_value,
+                    'experiment': 'test_experiment',
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
+            )
+        ]
     )
-    ])
 
-    workflow_mock.execute_activity_method.assert_has_calls([
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': True, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {
+                        'success': True,
+                        'version': 'test_version',
+                        'mlflow_run_id': 'test_mlflow_run_id',
+                        'mlflow_experiment_id': 'test_mlflow_experiment_id',
+                    },
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
             call(
                 Activities.export_data_to_postgres,
                 {
                     **metadata,
-                'data': workflow_mock.execute_activity_method.return_value,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
                     'schema': input_data['schema'],
                     'table_name': input_data['table_name'],
                 },
                 retry_policy=ANY,
-            start_to_close_timeout=ANY
+                start_to_close_timeout=ANY,
             )
-    ])
+        ]
+    )
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            {'success': False, 'object_key': 'test_object_key'},
+            {'success': True, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_called_once_with(
+        Activities.query_to_minio,
+        {
+            **metadata,
+            'query': input_data['query'],
+            'datetime_columns': input_data.get('datetime_columns', []),
+            'model_name': input_data['model_name'],
+            'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
+        },
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )
+
+    workflow_mock.execute_local_activity_method.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
+async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
+    input_data = {
+        'model_id': 'test_model_id',
+        'model_name': 'test_model',
+        'workflow_name': 'minimal_retrain',
+        'schedule_name': 'test_schedule',
+        'query': 'test_query',
+        'schema': 'test_schema',
+        'table_name': 'test_table',
+        'model_config': {
+            'target': 'test_target',
+            'transform_flavor': 'test_transform_flavor',
+            'predict_flavor': 'test_predict_flavor',
+        },
+    }
+
+    workflow_mock.execute_activity_method = AsyncMock(
+        side_effect=[
+            {'success': True, 'object_key': 'test_object_key'},
+            {'success': False, 'experiment': 'test_experiment'},
+            {
+                'success': True,
+                'version': 'test_version',
+                'mlflow_run_id': 'test_mlflow_run_id',
+                'mlflow_experiment_id': 'test_mlflow_experiment_id',
+            },
+            {'report': 'test_report'},
+        ]
+    )
+
+    await minimal_retrain.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.query_to_minio,
+                {
+                    **metadata,
+                    'query': input_data['query'],
+                    'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
+                    'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.retrain_model,
+                {
+                    **metadata,
+                    'object_key': 'test_object_key',
+                    'model_name': input_data['model_name'],
+                    'model_config': input_data['model_config'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_local_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.format_retrain_report,
+                {
+                    **metadata,
+                    'experiment_response': {'success': False, 'experiment': 'test_experiment'},
+                    'model_name': input_data['model_name'],
+                    'model_id': input_data['model_id'],
+                    'update_report': {},
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.export_data_to_postgres,
+                {
+                    **metadata,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
+                    'schema': input_data['schema'],
+                    'table_name': input_data['table_name'],
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
+    assert workflow_mock.execute_activity_method.call_count == 3
+    assert workflow_mock.execute_local_activity_method.call_count == 1

tests/laborious/workflows/test_predictions_batch.py

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

validate.sh

@@ -0,0 +1,99 @@
+#!/bin/bash
+# Model Manager Code Validation Script
+# This script runs all code quality checks before committing or deploying
+
+set -e  # Exit on any error
+
+# Colors for output
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+BLUE='\033[0;34m'
+NC='\033[0m' # No Color
+
+echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║  Model Manager - Code Validation Suite                ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
+echo ""
+
+# Check if virtual environment is activated
+if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then
+    echo -e "${YELLOW}⚠️  Warning: No virtual environment detected${NC}"
+    echo -e "${YELLOW}   Consider activating your venv/conda environment${NC}"
+    echo ""
+fi
+
+# Function to run a validation step
+run_step() {
+    local step_name=$1
+    local step_command=$2
+    
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    echo -e "${BLUE}▶ ${step_name}${NC}"
+    echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
+    
+    if eval "$step_command"; then
+        echo -e "${GREEN}✅ ${step_name} - PASSED${NC}"
+        echo ""
+        return 0
+    else
+        echo -e "${RED}❌ ${step_name} - FAILED${NC}"
+        echo ""
+        return 1
+    fi
+}
+
+# Track failures
+FAILED_STEPS=()
+
+# Step 1: Code Formatting Check (Ruff)
+if ! run_step "1. Code Formatting (Ruff)" "ruff format laborious/ tests/ && ruff format --check laborious/ tests/"; then
+    FAILED_STEPS+=("Code Formatting")
+fi
+
+# Step 2: Linting (Ruff)
+if ! run_step "2. Code Linting (Ruff)" "ruff check --fix laborious/ tests/"; then
+    FAILED_STEPS+=("Linting")
+fi
+
+# Step 3: Type Checking (mypy)
+if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then
+    FAILED_STEPS+=("Type Checking")
+fi
+
+# Step 4: Security Analysis (Bandit)
+if ! run_step "4. Security Analysis (Bandit)" "bandit -r laborious/ -ll -q"; then
+    FAILED_STEPS+=("Security Analysis")
+fi
+
+# Step 5: Unit Tests (pytest)
+if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=laborious --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then
+    FAILED_STEPS+=("Unit Tests")
+fi
+
+# Summary
+echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}"
+echo -e "${BLUE}║  Validation Summary                                    ║${NC}"
+echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}"
+echo ""
+
+if [ ${#FAILED_STEPS[@]} -eq 0 ]; then
+    echo -e "${GREEN}✅ All validation checks passed!${NC}"
+    echo -e "${GREEN}   Your code is ready for commit/deployment.${NC}"
+    echo ""
+    exit 0
+else
+    echo -e "${RED}❌ Validation failed for the following steps:${NC}"
+    for step in "${FAILED_STEPS[@]}"; do
+        echo -e "${RED}   • ${step}${NC}"
+    done
+    echo ""
+    echo -e "${YELLOW}💡 Tips:${NC}"
+    echo -e "${YELLOW}   • Run 'ruff format laborious/ tests/' to auto-fix formatting${NC}"
+    echo -e "${YELLOW}   • Run 'ruff check --fix laborious/ tests/' to auto-fix linting issues${NC}"
+    echo -e "${YELLOW}   • Review mypy errors and add type hints where needed${NC}"
+    echo -e "${YELLOW}   • Check bandit warnings for security issues${NC}"
+    echo -e "${YELLOW}   • Fix failing tests or improve test coverage${NC}"
+    echo ""
+    exit 1
+fi

values.yaml

@@ -3,7 +3,7 @@
 # Declare variables to be passed into your templates.
 
 # This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/
-replicaCount: 1
+replicaCount: 2
 
 # This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/
 image:
@@ -11,7 +11,7 @@ image:
   # This sets the pull policy for images.
   pullPolicy: Always
   # Overrides the image tag whose default is the chart appVersion.
-  tag: "0.0.2"
+  tag: "0.0.3"
 
 0# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
 imagePullSecrets:
@@ -151,7 +151,7 @@ env:
   - name: GITHUB_REPO_URL
     value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git"
   - name: GITHUB_BRANCH
-    value: SIENTIAPDE-1222-ajustar-a-library-para-fazer-o-download-do-courier
+    value: SIENTIAPDE-1231-ajustar-o-retreino-do-courier-no-laborious
   - name: PYTHON_APP
     value: "laborious.worker.worker"
 
@@ -213,6 +213,17 @@ env:
   - name: MONGODB_TTL_INDEX_HOURS
     value: "1"
 
+  - name: MINIO_ENDPOINT_URL
+    value: "http://minio.minio.svc.cluster.local:9000"
+  - name: MINIO_ACCESS_KEY
+    value: "admin"
+  - name: MINIO_SECRET_KEY
+    value: "FvcxOPX55j"
+  - name: MINIO_REGION_NAME
+    value: "sa-east-1"
+  - name: MINIO_DEFAULT_BUCKET
+    value: "sientia"
+
 ssh:
   enabled: true
   secretName: git-ssh-key-sientia-laborious-worker