Code import - branch 0.6.0
This commit is contained in:
35
.env.example
Normal file
35
.env.example
Normal file
@@ -0,0 +1,35 @@
|
||||
REDIS_HOST="redis-master.redis.svc.cluster.local"
|
||||
REDIS_PORT="6379"
|
||||
REDIS_USERNAME="redis_username"
|
||||
REDIS_PASSWORD="redis_password"
|
||||
|
||||
MONGODB_USERNAME="mongo_username"
|
||||
MONGODB_PASSWORD="mongo_password"
|
||||
MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
|
||||
MONGODB_DATABASE="sientia"
|
||||
MONGODB_TTL_INDEX_HOURS="1"
|
||||
|
||||
EMAIL_SENDER="aignosi@aignosi.com.br"
|
||||
EMAIL_SENDER_PASSWORD="smtp_password"
|
||||
EMAIL_SMTP_SERVER="smtp.gmail.com"
|
||||
EMAIL_SMTP_PORT="587"
|
||||
|
||||
POSTGRES_HOST="paradedb-rw.paradedb.svc.cluster.local"
|
||||
POSTGRES_PORT="5432"
|
||||
POSTGRES_USER="sientia"
|
||||
POSTGRES_PASSWORD="sientia"
|
||||
POSTGRES_DBNAME="sientia"
|
||||
POSTGRES_MIN_CONNECTIONS="10"
|
||||
POSTGRES_MAX_CONNECTIONS="40"
|
||||
|
||||
LOG_LEVEL="DEBUG"
|
||||
HTTP_METRICS_PORT="9090"
|
||||
PROJECT_NAME="sientia-orchestrator"
|
||||
|
||||
TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233"
|
||||
TEMPORAL_NAMESPACE="default"
|
||||
TEMPORAL_SCOUTER_NAMESPACE="scouter"
|
||||
TEMPORAL_LABORIOUS_NAMESPACE="laborious"
|
||||
|
||||
RUNTIME=local
|
||||
# ACTIVITY_EXECUTOR_MAX_WORKERS=200
|
||||
17
.github/workflows/quality-gate.yml
vendored
Normal file
17
.github/workflows/quality-gate.yml
vendored
Normal file
@@ -0,0 +1,17 @@
|
||||
name: Quality gate
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
types: [ opened, synchronize, reopened ]
|
||||
|
||||
jobs:
|
||||
quality-gate:
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-quality-gate.yml@main
|
||||
permissions: write-all
|
||||
with:
|
||||
project_name: 'orchestrator'
|
||||
repositories: 'sientia-dataops-library'
|
||||
requirements_file: 'requirements-local.txt'
|
||||
secrets: inherit
|
||||
16
.github/workflows/release.yml
vendored
Normal file
16
.github/workflows/release.yml
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
name: Create Release on Merge to Main
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [closed]
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
release:
|
||||
if: github.event.pull_request.merged == true
|
||||
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
|
||||
permissions: write-all
|
||||
with:
|
||||
project_name: 'orchestrator'
|
||||
secrets: inherit
|
||||
51
.gitignore
vendored
Normal file
51
.gitignore
vendored
Normal file
@@ -0,0 +1,51 @@
|
||||
# Ignorar volumes do Docker
|
||||
docker-compose.override.yml
|
||||
**/db_data/
|
||||
**/kafka-volume/
|
||||
**/zookeeper-volume/
|
||||
**/mage_data/
|
||||
**/minio_data/
|
||||
**/venv/
|
||||
**/certs/*.pem
|
||||
**/certs/*.der
|
||||
**/certs/*.csr
|
||||
**/deploy/*.yaml
|
||||
scouter/.file_versions/
|
||||
scouter/pipelines/**/triggers.yaml
|
||||
**/postgres_data/**
|
||||
**/redis_data/**
|
||||
# Ignorar arquivos e diretórios de cache do Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
*.pyd
|
||||
|
||||
# Ignorar logs
|
||||
*.log
|
||||
|
||||
# Ignorar arquivos de configuração locais
|
||||
.vscode/
|
||||
.pytest_cache/
|
||||
.idea/
|
||||
*.swp
|
||||
|
||||
# Ignorar arquivos temporários
|
||||
*.tmp
|
||||
*.bak
|
||||
*.old
|
||||
.secret
|
||||
|
||||
# Ignorar coverage
|
||||
htmlcov/
|
||||
.coverage
|
||||
coverage.xml
|
||||
|
||||
# git keys
|
||||
git_key*
|
||||
|
||||
git_log
|
||||
|
||||
.env
|
||||
|
||||
openspec/
|
||||
.cursor/
|
||||
902
README.md
Normal file
902
README.md
Normal file
@@ -0,0 +1,902 @@
|
||||
# SIENTIA DataOps Orchestrator Temporal
|
||||
|
||||
A high-performance, scalable workflow orchestration system built on Temporal.io for automated pipeline management, notification delivery, and resource coordination. The Orchestrator provides enterprise-grade workflow automation, real-time alerting, and comprehensive monitoring capabilities for the SIENTIA platform.
|
||||
|
||||
## 📑 Table of Contents
|
||||
|
||||
- [Features](#features)
|
||||
- [Core Functionality](#core-functionality)
|
||||
- [Advanced Capabilities](#advanced-capabilities)
|
||||
- [Architecture](#architecture)
|
||||
- [Architecture Principles](#architecture-principles)
|
||||
- [Task Queue Isolation](#2-task-queue-isolation)
|
||||
- [Workflows](#-workflows)
|
||||
- [Orchestrator Workflow](#1-orchestrator-workflow-orchestratorpy)
|
||||
- [Alerts Workflow](#2-alerts-workflow-alertspy)
|
||||
- [Reports Workflow](#3-reports-workflow-reportspy)
|
||||
- [Subworkflows](#subworkflows)
|
||||
- [Load Notification Package](#1-load-notification-package-load_notification_packagepy)
|
||||
- [Process Notifications](#2-process-notifications-process_notificationspy)
|
||||
- [Notification Filtering System](#-notification-filtering-system)
|
||||
- [Prerequisites](#-prerequisites)
|
||||
- [Installation](#-installation)
|
||||
- [How to Run](#-how-to-run)
|
||||
- [Configuration](#-configuration)
|
||||
- [Monitoring and Metrics](#-monitoring-and-metrics)
|
||||
- [Testing](#-testing)
|
||||
- [Code Quality & Validation](#-code-quality--validation)
|
||||
- [Development](#-development)
|
||||
- [Troubleshooting](#-troubleshooting)
|
||||
- [Performance Tuning](#-performance-tuning)
|
||||
- [Contributing](#-contributing)
|
||||
- [License](#-license)
|
||||
- [Support](#-support)
|
||||
|
||||
## Features
|
||||
|
||||
### Core Functionality
|
||||
- **Pipeline Orchestration**: Automated deployment and management of data processing pipelines
|
||||
- **Real-time Notifications**: Intelligent alert filtering and delivery with TTL management
|
||||
- **Resource Management**: Dynamic OPC server slot allocation and active ingestor monitoring
|
||||
- **Schedule Management**: Temporal-based workflow scheduling with automatic retry policies
|
||||
- **Multi-namespace Support**: Separate workflow queues for scouter and laborious operations
|
||||
|
||||
### Advanced Capabilities
|
||||
- **Incremental Processing**: Timestamp-based data loading to avoid reprocessing
|
||||
- **Intelligent Notification Filtering**: Advanced filtering system with:
|
||||
- User group-based notification filtering with custom policies
|
||||
- TTL-based duplicate prevention for alerts
|
||||
- Ignore lists for specific notifications
|
||||
- Persistent alert detection for ongoing issues
|
||||
- **Auto-scaling Workers**: Multiple worker instances with task queue isolation
|
||||
- **Comprehensive Logging**: Structured logging with PostgreSQL audit trails
|
||||
- **Prometheus Metrics**: Real-time monitoring and alerting integration
|
||||
- **Template-based Email Generation**: Jinja2-powered HTML email templates
|
||||
|
||||
## Architecture
|
||||
|
||||
The SIENTIA DataOps Orchestrator 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 data pipeline environments.
|
||||
|
||||
### Architecture Principles
|
||||
|
||||
#### 1. **Separation of Concerns**
|
||||
- **Worker Layer**: Manages Temporal workers, task queues, Prometheus metrics, and application lifecycle
|
||||
- **Workflow Layer**: Orchestrates business logic and process coordination with subworkflows for modularity
|
||||
- **Activity Layer**: Implements specific operations and external system interactions with error handling
|
||||
- **Data Layer**: Handles data persistence (MongoDB, PostgreSQL), caching (Redis), and external service connections
|
||||
|
||||
#### 2. **Task Queue Isolation**
|
||||
- **Orchestrator Queue**: Pipeline and resource management workflows (orchestrator workflow)
|
||||
- **Alerts Queue**: Real-time error notification workflows (alerts workflow, load_notification_package, process_notifications subworkflows)
|
||||
- **Reports Queue**: Scheduled reporting and summary workflows (reports workflow, load_notification_package, process_notifications subworkflows)
|
||||
|
||||
#### 3. **Fault Tolerance & Resilience**
|
||||
- **Automatic Retry Policies**: Configurable retry strategies for transient failures
|
||||
- **Graceful Degradation**: System continues operating with reduced functionality
|
||||
- **Comprehensive Error Handling**: Detailed error reporting and notification integration
|
||||
- **Connection Management**: Automatic reconnection for SMTP and database services
|
||||
|
||||
#### 4. **Scalability & Performance**
|
||||
- **Horizontal Scaling**: Multiple worker instances for load distribution
|
||||
- **Connection Pooling**: Optimized database and Redis connections
|
||||
- **Asynchronous Processing**: Non-blocking operations for improved throughput
|
||||
- **Resource Optimization**: Intelligent slot allocation and ingestor management
|
||||
|
||||
## 🔄 Workflows
|
||||
|
||||
### Main Workflows
|
||||
|
||||
#### 1. Orchestrator Workflow (`orchestrator.py`)
|
||||
|
||||
The **Orchestrator** workflow is the main coordination workflow that manages pipeline deployment and resource allocation across the SIENTIA platform.
|
||||
|
||||
**Purpose**:
|
||||
- **Pipeline Management**: Coordinates deployment of scouter and laborious pipelines
|
||||
- **Resource Allocation**: Manages OPC server slots and active ingestor distribution
|
||||
- **Schedule Synchronization**: Ensures Temporal schedules match MongoDB configurations
|
||||
- **Infrastructure Management**: Creates, updates, and deletes workflow schedules
|
||||
|
||||
**Execution Flow**:
|
||||
1. **Parallel Data Loading**: Concurrently loads pipelines, OPC servers, orchestrated schedules, OPC slots, and active ingestors
|
||||
2. **Parallel Processing**: Formats orchestrated schedules and processes new schedules and slots
|
||||
3. **Parallel Config Creation**: Creates schedule and slot action configurations, normalizes schedules, and creates collections with TTL indexes
|
||||
4. **Parallel Operations**: Executes slot deletions, slot updates, schedule deletions, schedule creations, and schedule updates concurrently
|
||||
5. **Parallel Reports & Timestamps**: Generates orchestration reports and updates MongoDB timestamps for created/updated/deleted pipelines
|
||||
|
||||
**Input Parameters**:
|
||||
```json
|
||||
{
|
||||
"schedule_name": "hourly_orchestration",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [
|
||||
{"$match": {"active": true}},
|
||||
{"$sort": {"updated_at": -1}}
|
||||
]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph "Parallel Data Loading"
|
||||
A[1. aggregate_documents_in_mongodb]
|
||||
B[2. find_documents_in_mongodb<br/>OPC Servers]
|
||||
C[3. find_documents_in_mongodb<br/>Orchestrated Schedules]
|
||||
D[4. load_opc_slots]
|
||||
E[5. load_active_ingestors]
|
||||
end
|
||||
|
||||
subgraph "Parallel Processing"
|
||||
F[6. format_schedule_config]
|
||||
G[7. process_schedules]
|
||||
H[8. process_slots]
|
||||
end
|
||||
|
||||
subgraph "Parallel Config Creation"
|
||||
I[9. create_schedule_config]
|
||||
J[10. create_slot_config]
|
||||
K[11. normalize_schedules]
|
||||
L[12. create_collection_with_ttl_index]
|
||||
end
|
||||
|
||||
subgraph "Parallel Operations"
|
||||
M[13. delete_slots]
|
||||
N[14. update_slots]
|
||||
O[15. delete_schedules]
|
||||
P[16. create_schedules]
|
||||
Q[17. update_schedules]
|
||||
end
|
||||
|
||||
subgraph "Parallel Reports & Timestamps"
|
||||
R[18. report_schedule_orchestration]
|
||||
S[19. report_slot_orchestration]
|
||||
T[20. update_pipelines_timestamps]
|
||||
U[21. create_pipelines_timestamps]
|
||||
V[22. delete_pipelines_timestamps]
|
||||
end
|
||||
|
||||
A --> F
|
||||
F --> I
|
||||
I --> M
|
||||
M --> R
|
||||
|
||||
A -.-> MongoDB1[(MongoDB)]
|
||||
D -.-> Redis1[(Redis)]
|
||||
K -.-> Temporal[(Temporal)]
|
||||
R -.-> Reports[Reports]
|
||||
```
|
||||
|
||||
Note: Green blocks (🟩) indicate parallel processing operations that run concurrently for improved performance.
|
||||
|
||||
|
||||
#### 2. Alerts Workflow (`alerts.py`)
|
||||
|
||||
The **Alerts** workflow processes and sends real-time error notifications to configured user groups with intelligent filtering and duplicate prevention.
|
||||
|
||||
**Purpose**:
|
||||
- **Error Alerting**: Immediate notification of ERROR-level events
|
||||
- **TTL Management**: Prevents alert spam using configurable time-to-live settings
|
||||
- **Group Filtering**: Sends alerts only to relevant user groups
|
||||
- **Persistent Monitoring**: Tracks and escalates persistent issues
|
||||
|
||||
**Execution Flow**:
|
||||
1. **Notification Loading**: Subworkflow loads ERROR-level notifications from MongoDB with timestamp filtering
|
||||
2. **Configuration Loading**: Loads active receiver group configurations from MongoDB
|
||||
3. **Alert Filtering**: Applies TTL-based filtering to identify core_alerts (new) and persistent_alerts (ongoing issues)
|
||||
4. **Group Filtering**: Filters notifications by user group content policies and ignore lists
|
||||
5. **Email Processing**: Subworkflow generates HTML emails and sends to receiver groups
|
||||
6. **Cache Storage**: Stores sent notification cache in Redis with configurable TTL to prevent duplicates
|
||||
7. **Audit Logging**: Logs delivery results to PostgreSQL for monitoring
|
||||
|
||||
**Input Parameters**:
|
||||
```json
|
||||
{
|
||||
"schedule_name": "error_alerts",
|
||||
"notification_ttl": 3600,
|
||||
"sent_ttl": 7200
|
||||
}
|
||||
```
|
||||
|
||||
#### Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. load_notification_package🔃] --> B[2. filter_notification_alerts] --> C[3. process_notifications🔃]
|
||||
|
||||
A -.-> MongoDB[(MongoDB)]
|
||||
A -.-> Redis[(Redis)]
|
||||
B -.-> Filters[Report Filters]
|
||||
C -.-> Email[(Email)]
|
||||
C -.-> PostgreSQL[(PostgreSQL)]
|
||||
|
||||
style A fill:#000,color:#fff
|
||||
style C fill:#000,color:#fff
|
||||
```
|
||||
|
||||
#### 3. Reports Workflow (`reports.py`)
|
||||
|
||||
The **Reports** workflow generates and sends scheduled comprehensive reports to configured user groups.
|
||||
|
||||
**Purpose**:
|
||||
- **Scheduled Reporting**: Regular summary reports of system activity
|
||||
- **Comprehensive Coverage**: Includes all notification levels (not just errors)
|
||||
- **Group Management**: Customizable reports per user group
|
||||
- **Audit Trail**: Complete logging of report delivery
|
||||
|
||||
**Execution Flow**:
|
||||
1. **Notification Loading**: Subworkflow loads all notifications from MongoDB (any level) with timestamp filtering
|
||||
2. **Configuration Loading**: Loads active receiver group configurations from MongoDB
|
||||
3. **Report Filtering**: Filters notifications by user group content policies (must include "reports") and ignore lists
|
||||
4. **Email Processing**: Subworkflow generates HTML reports organized by notification level and model, then sends to receiver groups
|
||||
5. **Audit Logging**: Logs delivery results to PostgreSQL for monitoring and tracking
|
||||
|
||||
**Input Parameters**:
|
||||
```json
|
||||
{
|
||||
"schedule_name": "reports",
|
||||
}
|
||||
```
|
||||
|
||||
#### Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. load_notification_package🔃] --> B[2. filter_notification_reports] --> C[3. process_notifications🔃]
|
||||
|
||||
A -.-> MongoDB[(MongoDB)]
|
||||
A -.-> Redis[(Redis)]
|
||||
B -.-> Filters[Report Filters]
|
||||
C -.-> Email[(Email)]
|
||||
C -.-> PostgreSQL[(PostgreSQL)]
|
||||
|
||||
style A fill:#000,color:#fff
|
||||
style C fill:#000,color:#fff
|
||||
```
|
||||
|
||||
### Subworkflows
|
||||
|
||||
#### 1. Load Notification Package (`load_notification_package.py`)
|
||||
|
||||
**Purpose**: Centralized notification data loading and configuration management for both alerts and reports workflows.
|
||||
|
||||
**Key Features**:
|
||||
- **Incremental Processing**: Uses Redis timestamps for efficient data loading
|
||||
- **Configuration Management**: Loads active receiver group configurations
|
||||
- **Data Validation**: Ensures complete data packages before processing
|
||||
- **Timestamp Management**: Updates last processed timestamps
|
||||
|
||||
**Input Parameters**:
|
||||
```json
|
||||
{
|
||||
"metadata": {"workflow_name": "alerts", "schedule_name": "error_alerts"},
|
||||
"mail_type": "Alerts",
|
||||
"base_data_filter": {"level": "ERROR"}
|
||||
}
|
||||
```
|
||||
|
||||
**Returns**:
|
||||
- `last_timestamp` (str | None): Last processed timestamp
|
||||
- `notification_package` (list[dict]): Retrieved notifications
|
||||
- `sending_configs` (list[dict]): Active receiver group configurations
|
||||
|
||||
**Key Activities**:
|
||||
- `get_last_data_timestamp`: Retrieves last processed timestamp from Redis for incremental processing
|
||||
- `find_documents_in_mongodb`: Loads active receiver group configurations from MongoDB
|
||||
- `load_latest_data`: Loads notifications with timestamp filtering from notification_queue collection
|
||||
- `put_last_data_timestamp`: Updates last processed timestamp in Redis with 5-hour TTL
|
||||
|
||||
**Architecture**:
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph "Parallel Loading"
|
||||
A[1. get_last_data_timestamp]
|
||||
B[2. find_documents_in_mongodb<br/>Receiver Groups]
|
||||
end
|
||||
|
||||
C[3. load_latest_data] --> D[4. put_last_data_timestamp]
|
||||
|
||||
A --> C
|
||||
B --> D
|
||||
|
||||
A -.-> Redis1[(Redis)]
|
||||
B -.-> MongoDB1[(MongoDB)]
|
||||
C -.-> MongoDB2[(MongoDB)]
|
||||
D -.-> Redis2[(Redis)]
|
||||
```
|
||||
|
||||
#### 2. Process Notifications (`process_notifications.py`)
|
||||
|
||||
**Purpose**: Handles email generation, delivery, and audit logging for notification workflows.
|
||||
|
||||
**Key Features**:
|
||||
- **HTML Generation**: Creates formatted email content for each receiver group using Jinja2 templates
|
||||
- **Email Delivery**: Sends emails with attachment support and error handling
|
||||
- **Audit Logging**: Records delivery status and metrics in PostgreSQL
|
||||
- **Error Recovery**: Handles SMTP failures with detailed error reporting
|
||||
- **Template Support**: Uses customizable HTML templates for different email types
|
||||
|
||||
**Input Parameters**:
|
||||
```json
|
||||
{
|
||||
"metadata": {"workflow_name": "alerts", "schedule_name": "error_alerts"},
|
||||
"mail_type": "Alerts",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "log_report",
|
||||
"notification_package": {
|
||||
"group_name": "admin_team",
|
||||
"members": ["admin@example.com"],
|
||||
"notifications": [...]
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Returns**:
|
||||
- `log_report` (list[dict]): Detailed delivery status for each notification
|
||||
|
||||
**Key Activities**:
|
||||
- `build_email_html`: Generates HTML content using Jinja2 templates organized by notification level and model
|
||||
- `send_email`: Delivers emails to receiver groups with attachment support and automatic SMTP reconnection
|
||||
- `format_log_report`: Formats delivery results into DataFrame structure for database storage, aggregating by notification ID and trigger
|
||||
- `export_data_to_postgres`: Stores audit logs in PostgreSQL with timestamp conversion
|
||||
|
||||
**Architecture**:
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. build_email_html] --> B[2. send_email] --> C[3. format_log_report] --> D[4. export_data_to_postgres]
|
||||
|
||||
A -.-> HTML[HTML Generator]
|
||||
B -.-> SMTP[(SMTP)]
|
||||
C -.-> Formatter[Log Formatter]
|
||||
D -.-> PostgreSQL[(PostgreSQL)]
|
||||
|
||||
```
|
||||
|
||||
## 🔔 Notification Filtering System
|
||||
|
||||
The orchestrator includes an advanced notification filtering system that prevents alert spam and ensures relevant notifications reach the appropriate user groups.
|
||||
|
||||
### **Alert Filtering (`filter_notification_alerts`)**
|
||||
- **Purpose**: Filters error-level notifications for immediate alerts
|
||||
- **TTL Management**: Prevents duplicate alerts using configurable time-to-live settings
|
||||
- **Persistent Detection**: Identifies ongoing issues that require escalation
|
||||
- **Group-based Filtering**: Routes notifications to appropriate receiver groups
|
||||
- **Ignore Lists**: Supports notification exclusion per group
|
||||
|
||||
### **Report Filtering (`filter_notification_reports`)**
|
||||
- **Purpose**: Filters all notification levels for comprehensive scheduled reports
|
||||
- **Comprehensive Coverage**: Includes ERROR, WARNING, INFO, and DEBUG levels
|
||||
- **Group Customization**: Applies different content policies per receiver group with ignore list support
|
||||
- **Scheduled Processing**: Designed for regular report generation without TTL-based duplicate prevention
|
||||
- **Duplicate Prevention**: Prevents duplicate notifications within the same report using trigger and notification ID keys
|
||||
|
||||
### **Notification Caching (`store_notification_cache`)**
|
||||
- **Purpose**: Manages Redis-based notification cache for TTL enforcement
|
||||
- **TTL Support**: Configurable expiration times for different notification types
|
||||
- **Duplicate Prevention**: Ensures notifications aren't sent repeatedly within TTL window
|
||||
- **Key Management**: Uses structured keys for efficient cache lookups
|
||||
|
||||
### Key Components
|
||||
|
||||
#### **Worker (`orchestrator/worker/worker.py`)**
|
||||
- **Purpose**: Main application orchestrator managing Temporal workers and task queues
|
||||
- **Responsibilities**:
|
||||
- Temporal client initialization and connection management with SDK metrics
|
||||
- Worker lifecycle management and graceful shutdown via `sientia_do.temporal.worker.prepare_worker`
|
||||
- Task queue configuration (orchestrator, alerts, reports) with dedicated workers on `<workflow>-<runtime>-queue`
|
||||
- Prometheus metrics server initialization on HTTP_METRICS_PORT
|
||||
- Temporal SDK metrics server initialization on HTTP_SDK_METRICS_PORT
|
||||
- Notification handler setup and configuration
|
||||
- **Key Features**:
|
||||
- Multi-queue worker management with three dedicated workers (orchestrator, alerts, reports)
|
||||
- `RUNTIME` env var (default `legacy`) passed to every `prepare_worker` call
|
||||
- Sync blocking activities run on `prepare_worker`'s `activity_executor` thread pool; workflows stay `async def`
|
||||
- Application health metrics (app_up gauge) for Kubernetes liveness/readiness probes
|
||||
- Graceful shutdown with cleanup procedures for all connections and exit code propagation to Kubernetes
|
||||
- Parallel worker execution using asyncio.gather
|
||||
|
||||
#### **Activities (`orchestrator/activities/`)**
|
||||
- **Activities**: Main activity orchestrator combining all operations (TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres)
|
||||
- **TemporalManager**: Temporal schedule CRUD operations across scouter and laborious namespaces with search attributes
|
||||
- **SlotManager**: Redis-based OPC slot and cache management with notification filtering and TTL-based duplicate prevention
|
||||
- **MongoDB**: Document operations, aggregations, timestamp management, and TTL index creation
|
||||
- **Email**: SMTP operations with HTML generation, attachment support, and automatic reconnection handling
|
||||
- **Formatters**: Configuration processing, slot distribution algorithms, notification filtering for reports, and schedule/slot orchestration reporting
|
||||
- **Postgres**: PostgreSQL operations for audit logging and data export (via sientia-dataops-library)
|
||||
|
||||
#### **Utilities (`orchestrator/utils/`)**
|
||||
- **Connectors Configuration**: Database and service configuration management from environment variables
|
||||
- **Email Builder**: HTML email template generation and formatting using Jinja2 templates with support for alerts and reports
|
||||
- **Orchestrator Functions**: Pipeline configuration transformation utilities supporting:
|
||||
- `scouter`: OPC UA data collection workflows
|
||||
- `pi_web_api_scouter`: PI Web API data collection workflows
|
||||
- `predictions_batch`: ML prediction workflows with OPC write-back and multi-stage filtering
|
||||
- `minimal_retrain`: Model retraining workflows with SQL queries
|
||||
- `drift`: Data drift detection workflows
|
||||
- `simple_metrics`: Model performance metrics computation workflows
|
||||
- **Converters**: Data type conversion and validation utilities including frequency parsing for Temporal schedules
|
||||
- **Templates**: HTML email templates for alerts and reports (email_template.html, general_template.html)
|
||||
|
||||
## 📋 Prerequisites
|
||||
|
||||
### System Requirements
|
||||
- Python 3.11+
|
||||
- Temporal server/cluster
|
||||
- Redis server
|
||||
- MongoDB server
|
||||
- PostgreSQL database
|
||||
- SMTP server access
|
||||
|
||||
**Note**: External dependencies must be available either through:
|
||||
- Kubernetes cluster deployment
|
||||
- Docker Compose setup
|
||||
- Cloud-managed services
|
||||
- Local installations
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
### Local Development Setup
|
||||
|
||||
1. **Clone the repository**
|
||||
```bash
|
||||
git clone <repository-url>
|
||||
cd sientia-dataops-orchestrator_temporal
|
||||
```
|
||||
|
||||
2. **Create virtual environment**
|
||||
```bash
|
||||
python3.11 -m venv venv
|
||||
source ./venv/bin/activate
|
||||
```
|
||||
|
||||
3. **Install dependencies**
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
4. **Configure environment variables**
|
||||
```bash
|
||||
# Set required environment variables for services
|
||||
export TEMPORAL_HOST=localhost:7233
|
||||
export REDIS_HOST=localhost
|
||||
export MONGODB_URL=localhost:27017
|
||||
export POSTGRES_HOST=localhost
|
||||
# ... additional configuration
|
||||
```
|
||||
|
||||
## 📦 How to Run
|
||||
|
||||
### Running the Orchestrator Application
|
||||
|
||||
Use the provided script to run the application locally:
|
||||
|
||||
```bash
|
||||
# Make script executable (first time only)
|
||||
chmod +x run_local.sh
|
||||
|
||||
# Run the application
|
||||
./run_local.sh
|
||||
```
|
||||
|
||||
The script will:
|
||||
- Activate the virtual environment
|
||||
- Load environment variables from `.env`
|
||||
- Start the orchestrator worker application
|
||||
|
||||
### Running Tests and Coverage
|
||||
|
||||
Use the provided script to run tests with coverage:
|
||||
|
||||
```bash
|
||||
# Make script executable (first time only)
|
||||
chmod +x run_coverage.sh
|
||||
|
||||
# Run tests with coverage
|
||||
./run_coverage.sh
|
||||
```
|
||||
|
||||
### Manual Application Execution
|
||||
|
||||
For manual execution without scripts:
|
||||
|
||||
```bash
|
||||
# Activate virtual environment
|
||||
source ./venv/bin/activate
|
||||
|
||||
# Start the orchestrator worker
|
||||
python -m orchestrator.worker.worker
|
||||
```
|
||||
|
||||
## ⚙️ Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Description | Default | Required |
|
||||
|----------|-------------|---------|----------|
|
||||
| `TEMPORAL_HOST` | Temporal server address | `localhost:7233` | Yes |
|
||||
| `TEMPORAL_NAMESPACE` | Default Temporal namespace | `default` | No |
|
||||
| `TEMPORAL_SCOUTER_NAMESPACE` | Scouter workflow namespace | `scouter` | No |
|
||||
| `TEMPORAL_LABORIOUS_NAMESPACE` | Laborious workflow namespace | `laborious` | No |
|
||||
| `RUNTIME` | Runtime slice for orchestrator worker task queues (`orchestrator-<runtime>-queue`, etc.) | `legacy` | No |
|
||||
| `ACTIVITY_EXECUTOR_MAX_WORKERS` | Thread pool size for sync activities (all workers) | `200` | No |
|
||||
| `ORCHESTRATOR_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Orchestrator worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
|
||||
| `ALERTS_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Alerts worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
|
||||
| `REPORTS_ACTIVITY_EXECUTOR_MAX_WORKERS` | Per-worker override for the Reports worker | falls back to `ACTIVITY_EXECUTOR_MAX_WORKERS` | No |
|
||||
| `REDIS_HOST` | Redis server hostname | `localhost` | Yes |
|
||||
| `REDIS_PORT` | Redis server port | `6379` | Yes |
|
||||
| `REDIS_USERNAME` | Redis username | `default` | Yes |
|
||||
| `REDIS_PASSWORD` | Redis password | - | Yes |
|
||||
| `MONGODB_URL` | MongoDB server URL | `localhost:27017` | Yes |
|
||||
| `MONGODB_USERNAME` | MongoDB username | `root` | Yes |
|
||||
| `MONGODB_PASSWORD` | MongoDB password | - | Yes |
|
||||
| `MONGODB_DATABASE_NAME` | MongoDB database name | `sientia` | Yes |
|
||||
| `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes |
|
||||
| `POSTGRES_PORT` | PostgreSQL port | `5432` | Yes |
|
||||
| `POSTGRES_USER` | PostgreSQL username | `sientia` | Yes |
|
||||
| `POSTGRES_PASSWORD` | PostgreSQL password | - | Yes |
|
||||
| `POSTGRES_DBNAME` | PostgreSQL database | `sientia` | Yes |
|
||||
| `EMAIL_SENDER` | Sender email address | - | Yes |
|
||||
| `EMAIL_SENDER_PASSWORD` | SMTP password | - | Yes |
|
||||
| `EMAIL_SMTP_SERVER` | SMTP server | `smtp.gmail.com` | No |
|
||||
| `EMAIL_SMTP_PORT` | SMTP port | `587` | No |
|
||||
| `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No |
|
||||
| `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No |
|
||||
| `POSTGRES_MIN_CONNECTIONS` | PostgreSQL minimum connections | `10` | No |
|
||||
| `POSTGRES_MAX_CONNECTIONS` | PostgreSQL maximum connections | `40` | No |
|
||||
| `MONGODB_TTL_INDEX_HOURS` | MongoDB TTL index expiration in hours | `1` | No |
|
||||
| `PROJECT_NAME` | Project name for metrics and logging | `sientia-orchestrator` | No |
|
||||
| `LOG_LEVEL` | Application logging level | `INFO` | No |
|
||||
| `KAFKA_BOOTSTRAP_SERVERS` | Kafka bootstrap servers | - | No |
|
||||
|
||||
### Workflow Configuration
|
||||
|
||||
Temporal input configuration sample:
|
||||
|
||||
#### Orchestrator Workflow
|
||||
|
||||
```json
|
||||
{
|
||||
"schedule_name": "orchestrator-test",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [
|
||||
{
|
||||
"$lookup": {
|
||||
"from": "models",
|
||||
"localField": "model_id",
|
||||
"foreignField": "id",
|
||||
"as": "model_docs"
|
||||
}
|
||||
},
|
||||
{
|
||||
"$match": {
|
||||
"active": True
|
||||
}
|
||||
},
|
||||
{
|
||||
"$addFields": {
|
||||
"models": {
|
||||
"$arrayElemAt": [
|
||||
"$model_docs",
|
||||
0
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"$match": {
|
||||
"models.active": True
|
||||
}
|
||||
},
|
||||
{
|
||||
"$project": {
|
||||
"model_docs": 0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc-servers",
|
||||
"filters": {
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Alerts Workflow
|
||||
|
||||
```json
|
||||
{
|
||||
"schedule_name": "alerts",
|
||||
"notification_ttl": 5*60,
|
||||
"sent_ttl": 10*60
|
||||
}
|
||||
```
|
||||
|
||||
#### Reports Workflow
|
||||
|
||||
```json
|
||||
{
|
||||
"schedule_name": "reports",
|
||||
}
|
||||
```
|
||||
|
||||
## 📊 Monitoring and Metrics
|
||||
|
||||
The Orchestrator system exposes comprehensive Prometheus metrics:
|
||||
|
||||
### Application Metrics
|
||||
- `app_up`: Application health status gauge (1=healthy, 0=unhealthy) labeled by pod_id
|
||||
- `email_sent_count`: Email delivery counter labeled by pod_id, model_name, pipeline_name, and email_group
|
||||
|
||||
### Workflow Metrics
|
||||
- Schedule creation, update, and deletion success rates (via notification reports)
|
||||
- Notification processing times and error rates (via PostgreSQL audit logs)
|
||||
- Resource allocation and slot management metrics (via notification reports)
|
||||
- Temporal SDK metrics exposed on HTTP_SDK_METRICS_PORT (default: 9091)
|
||||
|
||||
### Database Metrics
|
||||
- MongoDB query performance and connection health (via sientia-dataops-library)
|
||||
- Redis operation counts and response times (via sientia-dataops-library)
|
||||
- PostgreSQL export operations and audit log metrics (via sientia-dataops-library)
|
||||
|
||||
## 🧪 Testing
|
||||
|
||||
### Test Structure
|
||||
```
|
||||
tests/
|
||||
├── orchestrator/ # Orchestrator workflow tests
|
||||
│ ├── test_activities.py
|
||||
│ ├── test_email.py
|
||||
│ ├── test_formatters.py
|
||||
│ ├── test_mongo_db.py
|
||||
│ ├── test_slot_manager.py
|
||||
│ ├── test_temporal_manager.py
|
||||
│ └── test_workflows.py
|
||||
├── activities/ # Activity implementation tests
|
||||
├── utils/ # Utility function tests
|
||||
└── integration/ # End-to-end workflow tests
|
||||
```
|
||||
|
||||
### Test Coverage
|
||||
The project maintains comprehensive test coverage including:
|
||||
- **Activity Tests**: Unit tests for all activity classes
|
||||
- **Workflow Tests**: Integration tests for workflow orchestration
|
||||
- **Utility Tests**: Tests for configuration builders and converters
|
||||
- **Database Tests**: Tests for MongoDB, Redis, and PostgreSQL operations
|
||||
- **Email Tests**: Tests for email generation and delivery
|
||||
- **Notification Tests**: Tests for filtering and caching logic
|
||||
|
||||
### Test Execution
|
||||
```bash
|
||||
# Install test dependencies
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
# Run unit tests with coverage (default testpaths=tests; E2E excluded)
|
||||
pytest --cov=orchestrator --cov-report=html
|
||||
|
||||
# Run specific test modules
|
||||
pytest tests/orchestrator/activities/test_mongo_db.py
|
||||
pytest tests/orchestrator/workflows/test_orchestrator.py
|
||||
```
|
||||
|
||||
### End-to-end tests
|
||||
|
||||
E2E tests live in `e2e/` and require **Docker** (testcontainers). They are **not** collected by default `pytest` at the repo root.
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
pip install -r requirements-dev.txt
|
||||
pytest e2e/ --override-ini testpaths=e2e -m e2e -v
|
||||
```
|
||||
|
||||
Coverage is kept separate from unit tests: set `COVERAGE_FILE=.coverage.e2e` when measuring E2E coverage (see `e2e/README.md`).
|
||||
|
||||
## 🛡️ Code Quality & Validation
|
||||
|
||||
### Overview
|
||||
|
||||
Since Python is not compiled, this project ships a validation workflow to catch issues early. Use the `validate.sh` script to run formatting, linting, type checks, security analysis, and tests in one command.
|
||||
|
||||
### Validation Tools
|
||||
|
||||
- **Ruff**: formatting and linting (fast, replaces Black/Flake8)
|
||||
- **mypy**: static typing checks
|
||||
- **Bandit**: security static analysis
|
||||
- **pytest**: unit/integration tests with coverage
|
||||
|
||||
### Tools Installation
|
||||
|
||||
```bash
|
||||
pip install -r requirements-dev.txt
|
||||
```
|
||||
|
||||
### Complete Validation (recommended)
|
||||
|
||||
```bash
|
||||
./validate.sh
|
||||
```
|
||||
|
||||
What `validate.sh` does:
|
||||
1. Checks formatting with Ruff
|
||||
2. Lints code with Ruff
|
||||
3. Runs mypy type checking
|
||||
4. Runs Bandit security analysis
|
||||
5. Executes pytest with coverage (generates HTML report)
|
||||
|
||||
Exit codes are propagated so CI can fail fast when quality gates are not met.
|
||||
|
||||
### Individual Commands
|
||||
|
||||
```bash
|
||||
# 1) Format check
|
||||
ruff format --check orchestrator/ tests/
|
||||
|
||||
# 2) Lint
|
||||
ruff check orchestrator/ tests/
|
||||
|
||||
# 3) Type check
|
||||
mypy orchestrator/
|
||||
|
||||
# 4) Security
|
||||
bandit -r orchestrator/ -ll
|
||||
|
||||
# 5) Tests with coverage
|
||||
pytest tests/ --cov=orchestrator --cov-report=html
|
||||
```
|
||||
|
||||
### Automatic Fixes
|
||||
|
||||
```bash
|
||||
# Apply formatting
|
||||
ruff format orchestrator/ tests/
|
||||
|
||||
# Autofix common lint issues
|
||||
ruff check --fix orchestrator/ tests/
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
Tooling is configured in `pyproject.toml` (lint rules, formatting, typing). Adjust thresholds and rules there as needed.
|
||||
|
||||
## 🔧 Development
|
||||
|
||||
### Project Structure
|
||||
```
|
||||
orchestrator/
|
||||
├── activities/ # Temporal activity implementations
|
||||
│ ├── activities.py # Main activities orchestrator
|
||||
│ ├── temporal_manager.py # Temporal schedule operations
|
||||
│ ├── slot_manager.py # Redis slot management and notification filtering
|
||||
│ ├── mongo_db.py # MongoDB operations
|
||||
│ ├── email.py # Email service operations
|
||||
│ ├── formatters.py # Configuration formatting and report filtering
|
||||
├── workflows/ # Temporal workflow definitions
|
||||
│ ├── orchestrator.py # Main orchestration workflow
|
||||
│ ├── alerts.py # Error alert workflow
|
||||
│ ├── reports.py # Scheduled report workflow
|
||||
│ └── subworkflows/ # Sub-workflow implementations
|
||||
│ ├── load_notification_package.py # Notification data loading
|
||||
│ └── process_notifications.py # Email processing and delivery
|
||||
├── worker/ # Worker implementation
|
||||
│ └── worker.py # Main worker orchestrator
|
||||
├── utils/ # Utility functions
|
||||
│ ├── connectors_config.py # Database configuration
|
||||
│ ├── email_builder.py # Email template generation
|
||||
│ ├── orchestrator_functions.py # Pipeline utilities
|
||||
│ ├── converters.py # Data type conversion utilities
|
||||
│ └── templates/ # Email HTML templates
|
||||
│ ├── email_template.html # Report email template
|
||||
│ └── general_template.html # General email template
|
||||
└── metrics.py # Prometheus metrics definitions
|
||||
```
|
||||
|
||||
### Adding New Features
|
||||
|
||||
1. **Follow Temporal patterns** for new workflows and activities
|
||||
2. **Add comprehensive docstrings** for all public methods
|
||||
3. **Include Prometheus metrics** for monitoring
|
||||
4. **Add unit tests** for new functionality
|
||||
5. **Update this README** with new features and configuration
|
||||
|
||||
## 🐛 Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
1. **Temporal Connection Failures**
|
||||
- Verify Temporal server is running and accessible
|
||||
- Check namespace configuration and permissions
|
||||
- Review server logs for connection issues
|
||||
|
||||
2. **Database Connection Issues**
|
||||
- Verify all database services are running
|
||||
- Check connection credentials and network access
|
||||
- Ensure proper connection pool configuration
|
||||
|
||||
3. **Email Delivery Failures**
|
||||
- Verify SMTP server configuration and credentials
|
||||
- Check email sender permissions and authentication
|
||||
- Review email delivery logs for specific errors
|
||||
|
||||
4. **Workflow Execution Failures**
|
||||
- Review activity error logs and notifications
|
||||
- Check MongoDB collection configurations
|
||||
- Verify input data format and required fields
|
||||
|
||||
5. **E2E / testcontainers leftovers**
|
||||
- If a run is interrupted, containers may keep running. List and remove them:
|
||||
```bash
|
||||
docker ps -a --filter label=org.testcontainers=true
|
||||
docker rm -f $(docker ps -aq --filter label=org.testcontainers=true)
|
||||
```
|
||||
|
||||
### Debug Mode
|
||||
|
||||
Enable debug logging by setting the log level:
|
||||
```bash
|
||||
export LOG_LEVEL=DEBUG
|
||||
```
|
||||
|
||||
## ⚡ Performance Tuning
|
||||
|
||||
### Key Parameters
|
||||
|
||||
- **Worker Concurrency**: Configure worker task limits in Temporal client
|
||||
- **Connection Pools**: Optimize database connection pool sizes
|
||||
- **Redis TTL**: Adjust cache TTL settings based on requirements
|
||||
- **Batch Sizes**: Configure notification processing batch sizes
|
||||
|
||||
### Scaling Considerations
|
||||
|
||||
- **Horizontal Scaling**: Deploy multiple worker instances
|
||||
- **Task Queue Distribution**: Use dedicated queues for different workflows
|
||||
- **Database Performance**: Optimize indexes and connection pooling
|
||||
- **Memory Management**: Monitor and configure appropriate resource limits
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
1. Fork the repository
|
||||
2. Create a feature branch
|
||||
3. Make your changes with comprehensive testing
|
||||
4. Update documentation and docstrings
|
||||
5. Submit a pull request
|
||||
|
||||
### Code Quality Standards
|
||||
|
||||
- Follow PEP 8 style guidelines
|
||||
- Include comprehensive docstrings for all public methods
|
||||
- Maintain test coverage above 80%
|
||||
- Use type hints where appropriate
|
||||
- Follow Temporal.io best practices
|
||||
|
||||
## 📄 License
|
||||
|
||||
This project is licensed under the terms specified in the LICENSE file.
|
||||
|
||||
## 🆘 Support
|
||||
|
||||
For support and questions:
|
||||
- Check the troubleshooting section above
|
||||
- Review the metrics and logs for error patterns
|
||||
- Open an issue in the project repository
|
||||
- Contact the development team
|
||||
|
||||
---
|
||||
|
||||
### Runtime queues and migration
|
||||
|
||||
- **Queue naming**: Downstream pipeline schedules target `<workflow_type>-<runtime>-queue` (for example `scouter-legacy-queue`). The orchestrator's own workers register on `orchestrator-<runtime>-queue`, `alerts-<runtime>-queue`, and `reports-<runtime>-queue`.
|
||||
- **Pipeline `runtime`**: Each pipeline JSON may include `runtime` (default `legacy` via `common_config`). Schedules without `runtime` use `legacy`.
|
||||
- **Schedule updates**: Changing a pipeline's `runtime` does not update an existing schedule's `task_queue` in place; the next orchestrator tick deletes and recreates the schedule on the new queue.
|
||||
- **Rollout order**: Deploy scouter/laborious worker fleets bound to the new `<workflow_type>-<runtime>-queue` family before redeploying the orchestrator. On the first tick after upgrade, orphan schedules on old queue names are normalized away and schedules are recreated on the new queues (no in-place rename).
|
||||
|
||||
**Note**: The SIENTIA DataOps Orchestrator is designed for production use in enterprise data environments. Ensure proper security configuration and network isolation for production deployments.
|
||||
35
e2e/README.md
Normal file
35
e2e/README.md
Normal file
@@ -0,0 +1,35 @@
|
||||
# Orchestrator end-to-end tests
|
||||
|
||||
End-to-end tests run every external dependency for real (MongoDB, Redis, PostgreSQL via testcontainers; SMTP via in-process `aiosmtpd`; Temporal via `WorkflowEnvironment.start_local()`).
|
||||
|
||||
## Requirements
|
||||
|
||||
- Docker (for testcontainers)
|
||||
- Python dev dependencies: `pip install -r requirements-dev.txt`
|
||||
|
||||
## Run locally
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
pytest e2e/ --override-ini testpaths=e2e -m e2e -v
|
||||
```
|
||||
|
||||
Stop on first failure:
|
||||
|
||||
```bash
|
||||
pytest e2e/ --override-ini testpaths=e2e -m e2e -x
|
||||
```
|
||||
|
||||
## Coverage (separate from unit tests)
|
||||
|
||||
```bash
|
||||
COVERAGE_FILE=.coverage.e2e pytest e2e/ --override-ini testpaths=e2e -m e2e --cov=orchestrator --cov-branch
|
||||
coverage combine .coverage .coverage.e2e
|
||||
coverage report
|
||||
```
|
||||
|
||||
Unit tests keep the default `.coverage` file; the E2E run must set `COVERAGE_FILE=.coverage.e2e` so reports do not overwrite each other.
|
||||
|
||||
## Scenario catalog
|
||||
|
||||
See [scenarios.md](scenarios.md) for numbered scenarios and which test module implements each case.
|
||||
0
e2e/__init__.py
Normal file
0
e2e/__init__.py
Normal file
446
e2e/conftest.py
Normal file
446
e2e/conftest.py
Normal file
@@ -0,0 +1,446 @@
|
||||
"""Pytest configuration and fixtures for orchestrator E2E tests."""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from google.protobuf.duration_pb2 import Duration
|
||||
from pymongo import MongoClient
|
||||
from redis import Redis
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.temporal.worker.prepare_worker import build_queue_name
|
||||
from sqlalchemy import create_engine
|
||||
from temporalio.api.enums.v1 import IndexedValueType
|
||||
from temporalio.api.operatorservice.v1 import AddSearchAttributesRequest
|
||||
from temporalio.api.workflowservice.v1 import (
|
||||
DescribeNamespaceRequest,
|
||||
RegisterNamespaceRequest,
|
||||
)
|
||||
from temporalio.client import Client
|
||||
from temporalio.common import SearchAttributeKey
|
||||
from temporalio.service import RPCError, RPCStatusCode
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
from testcontainers.core.container import DockerContainer
|
||||
from testcontainers.postgres import PostgresContainer
|
||||
|
||||
from e2e.helpers import MONGO_COLLECTIONS, ORCHESTRATOR_TASK_QUEUE
|
||||
from e2e.smtp_test_server import SmtpTestServer
|
||||
from e2e.stub_workflows import STUB_WORKFLOW_CLASSES
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.activities.formatters import schedule_types
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
from orchestrator.workflows.reports import Reports
|
||||
from orchestrator.workflows.subworkflows.load_notification_package import (
|
||||
LoadNotificationPackage,
|
||||
)
|
||||
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
|
||||
|
||||
DB_SCHEMA_SQL_PATH = Path(__file__).parent / 'db_schema.sql'
|
||||
E2E_DATABASE = 'orchestrator_test'
|
||||
E2E_RUNTIMES = ('legacy', 'gpu')
|
||||
MANAGED_NAMESPACES = ('scouter', 'laborious')
|
||||
|
||||
E2E_SEARCH_ATTRIBUTES = [
|
||||
SearchAttributeKey.for_keyword('model_id'),
|
||||
SearchAttributeKey.for_keyword('model_name'),
|
||||
SearchAttributeKey.for_keyword('orchestrated'),
|
||||
]
|
||||
|
||||
E2E_NAMESPACE_SEARCH_ATTRIBUTES = {
|
||||
'model_id': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
'model_name': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
'orchestrated': IndexedValueType.INDEXED_VALUE_TYPE_KEYWORD,
|
||||
}
|
||||
|
||||
|
||||
async def register_namespace_if_missing(env: WorkflowEnvironment, namespace: str) -> None:
|
||||
"""
|
||||
Register a Temporal namespace on the local dev server and wait until it is ready.
|
||||
|
||||
Args:
|
||||
env: Session WorkflowEnvironment from start_local().
|
||||
namespace: Namespace name to register.
|
||||
"""
|
||||
service = env.client.service_client
|
||||
try:
|
||||
await service.workflow_service.register_namespace(
|
||||
RegisterNamespaceRequest(
|
||||
namespace=namespace,
|
||||
workflow_execution_retention_period=Duration(seconds=86400),
|
||||
)
|
||||
)
|
||||
except RPCError as err:
|
||||
if err.status != RPCStatusCode.ALREADY_EXISTS:
|
||||
raise
|
||||
|
||||
deadline = time.monotonic() + 5.0
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
await service.workflow_service.describe_namespace(
|
||||
DescribeNamespaceRequest(namespace=namespace)
|
||||
)
|
||||
return
|
||||
except RPCError:
|
||||
await asyncio.sleep(0.1)
|
||||
raise TimeoutError(f'Namespace {namespace} not ready within 5s')
|
||||
|
||||
|
||||
async def ensure_namespace_search_attributes(
|
||||
env: WorkflowEnvironment, namespace: str
|
||||
) -> None:
|
||||
"""
|
||||
Register the orchestrator search attributes on a namespace, if missing.
|
||||
|
||||
The local Temporal dev server only registers search attributes on the default
|
||||
namespace at start time. Schedules created in additional namespaces fail with
|
||||
"no mapping defined for search attribute ..." unless we explicitly add the
|
||||
same attribute mappings to those namespaces via the operator service.
|
||||
|
||||
Args:
|
||||
env: Session WorkflowEnvironment from start_local().
|
||||
namespace: Namespace where attributes must be available.
|
||||
"""
|
||||
service = env.client.service_client
|
||||
try:
|
||||
await service.operator_service.add_search_attributes(
|
||||
AddSearchAttributesRequest(
|
||||
namespace=namespace,
|
||||
search_attributes=dict(E2E_NAMESPACE_SEARCH_ATTRIBUTES),
|
||||
)
|
||||
)
|
||||
except RPCError as err:
|
||||
if err.status != RPCStatusCode.ALREADY_EXISTS:
|
||||
raise
|
||||
|
||||
|
||||
def temporal_host_from_env(env: WorkflowEnvironment) -> str:
|
||||
"""Return target host:port for the in-process Temporal dev server."""
|
||||
return env.client.service_client.config.target_host
|
||||
|
||||
|
||||
def mongo_uri_from_container(mongo_container) -> str:
|
||||
"""Build a Mongo connection string for the testcontainer."""
|
||||
port = mongo_container.get_exposed_port(27017)
|
||||
return f'mongodb://localhost:{port}'
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def postgres_container():
|
||||
"""PostgreSQL testcontainer used by all E2E tests."""
|
||||
postgres = PostgresContainer('postgres:15')
|
||||
postgres.start()
|
||||
yield postgres
|
||||
postgres.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def mongo_container():
|
||||
"""MongoDB testcontainer used by real CoreNotificationHandler."""
|
||||
mongo = DockerContainer('mongo:7').with_exposed_ports(27017)
|
||||
mongo.start()
|
||||
yield mongo
|
||||
mongo.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
def redis_container():
|
||||
"""Redis testcontainer for slot and notification timestamp paths."""
|
||||
redis = DockerContainer('redis:7').with_exposed_ports(6379)
|
||||
redis.start()
|
||||
yield redis
|
||||
redis.stop()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def postgres_engine(postgres_container):
|
||||
"""SQLAlchemy engine bound to the PostgreSQL testcontainer."""
|
||||
engine = create_engine(postgres_container.get_connection_url())
|
||||
yield engine
|
||||
engine.dispose()
|
||||
|
||||
|
||||
def _create_schema_and_tables(engine):
|
||||
sql_text = DB_SCHEMA_SQL_PATH.read_text(encoding='utf-8')
|
||||
with engine.begin() as conn:
|
||||
conn.exec_driver_sql(sql_text)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def setup_postgres_schema_and_tables(postgres_engine):
|
||||
"""Recreate Postgres schema from e2e/db_schema.sql before each test."""
|
||||
_create_schema_and_tables(postgres_engine)
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mongo_uri(mongo_container):
|
||||
return mongo_uri_from_container(mongo_container)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def reset_mongo_collections(mongo_uri):
|
||||
"""Drop orchestrator-managed Mongo collections between tests."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
db = client[E2E_DATABASE]
|
||||
for name in MONGO_COLLECTIONS:
|
||||
db[name].drop()
|
||||
finally:
|
||||
client.close()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_client(redis_container):
|
||||
"""Redis client bound to the testcontainer."""
|
||||
port = int(redis_container.get_exposed_port(6379))
|
||||
client = Redis(host='localhost', port=port, decode_responses=True)
|
||||
yield client
|
||||
client.close()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
def reset_redis(redis_client):
|
||||
"""Flush Redis between tests."""
|
||||
redis_client.flushdb()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def smtp_server():
|
||||
"""Session-scoped in-process SMTP server."""
|
||||
server = SmtpTestServer()
|
||||
server.start()
|
||||
yield server
|
||||
server.stop()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def smtp_messages_cleanup(smtp_server):
|
||||
"""Clear captured SMTP messages between tests."""
|
||||
smtp_server.clear()
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope='session')
|
||||
async def temporal_env():
|
||||
"""Real Temporal dev server (schedule APIs supported)."""
|
||||
env = await WorkflowEnvironment.start_local(search_attributes=E2E_SEARCH_ATTRIBUTES)
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
await register_namespace_if_missing(env, namespace)
|
||||
await ensure_namespace_search_attributes(env, namespace)
|
||||
yield env
|
||||
await env.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture(scope='session')
|
||||
def temporal_host(temporal_env):
|
||||
return temporal_host_from_env(temporal_env)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(autouse=True)
|
||||
async def cleanup_temporal_schedules(temporal_env):
|
||||
"""Delete orphan schedules in scouter/laborious before each test."""
|
||||
host = temporal_host_from_env(temporal_env)
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
client = await Client.connect(host, namespace=namespace)
|
||||
async for schedule in await client.list_schedules():
|
||||
handle = client.get_schedule_handle(schedule.id)
|
||||
await handle.delete()
|
||||
yield
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
def mock_logger():
|
||||
"""Logger double with readable console output for E2E runs."""
|
||||
logger = MagicMock(spec=Logger)
|
||||
logger.info = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.debug = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.error = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.warning = MagicMock(side_effect=lambda msg: print(f'[LOG] {msg}'))
|
||||
logger.custom_info = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_debug = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_error = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
logger.custom_warning = MagicMock(side_effect=lambda msg, _meta=None: print(f'[LOG] {msg}'))
|
||||
return logger
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def notification_handler(mock_logger, mongo_container):
|
||||
"""Real notification handler using MongoDB testcontainer."""
|
||||
handler = CoreNotificationHandler(
|
||||
connection_string=mongo_uri_from_container(mongo_container),
|
||||
database=E2E_DATABASE,
|
||||
logger=mock_logger,
|
||||
project_name='orchestrator-e2e',
|
||||
)
|
||||
try:
|
||||
yield handler
|
||||
finally:
|
||||
handler.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def notification_inserts(notification_handler):
|
||||
"""Spy on real Mongo insert calls issued by notification handler."""
|
||||
collection = notification_handler.mongo_collection
|
||||
original_insert_one = collection.insert_one
|
||||
spy = MagicMock(wraps=original_insert_one)
|
||||
collection.insert_one = spy
|
||||
try:
|
||||
yield spy
|
||||
finally:
|
||||
collection.insert_one = original_insert_one
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def test_activities(
|
||||
postgres_container,
|
||||
mongo_container,
|
||||
redis_container,
|
||||
smtp_server,
|
||||
mock_logger,
|
||||
notification_handler,
|
||||
temporal_host,
|
||||
):
|
||||
"""Real Activities wired to testcontainers and in-process SMTP."""
|
||||
mongo_port = mongo_container.get_exposed_port(27017)
|
||||
redis_port = int(redis_container.get_exposed_port(6379))
|
||||
|
||||
activities = Activities(
|
||||
temporal_config={
|
||||
'temporal_host': temporal_host,
|
||||
'temporal_scouter_namespace': 'scouter',
|
||||
'temporal_laborious_namespace': 'laborious',
|
||||
},
|
||||
redis_config={
|
||||
'host': 'localhost',
|
||||
'port': redis_port,
|
||||
'username': '',
|
||||
'password': '',
|
||||
},
|
||||
mongodb_config={
|
||||
'connection_string': f'mongodb://localhost:{mongo_port}',
|
||||
'database_name': E2E_DATABASE,
|
||||
'ttl_index_seconds': 3600,
|
||||
},
|
||||
email_config={
|
||||
'sender_email': 'e2e@example.com',
|
||||
'sender_password': '',
|
||||
'smtp_server': smtp_server.host,
|
||||
'smtp_port': smtp_server.port,
|
||||
},
|
||||
postgres_config={
|
||||
'host': 'localhost',
|
||||
'port': int(postgres_container.get_exposed_port(5432)),
|
||||
'user': postgres_container.username,
|
||||
'password': postgres_container.password,
|
||||
'dbname': postgres_container.dbname,
|
||||
'min_connections': 1,
|
||||
'max_connections': 5,
|
||||
},
|
||||
logger=mock_logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
await activities.connect_to_temporal()
|
||||
try:
|
||||
yield activities
|
||||
finally:
|
||||
activities.shutdown()
|
||||
|
||||
|
||||
def _orchestrator_activity_list(activities: Activities) -> list:
|
||||
return [
|
||||
activities.load_active_ingestors,
|
||||
activities.load_opc_slots,
|
||||
activities.update_slots,
|
||||
activities.delete_slots,
|
||||
activities.aggregate_documents_in_mongodb,
|
||||
activities.find_documents_in_mongodb,
|
||||
activities.update_pipelines_timestamps,
|
||||
activities.create_pipelines_timestamps,
|
||||
activities.delete_pipelines_timestamps,
|
||||
activities.create_collection_with_ttl_index,
|
||||
activities.create_schedules,
|
||||
activities.update_schedules,
|
||||
activities.delete_schedules,
|
||||
activities.normalize_schedules,
|
||||
activities.process_schedules,
|
||||
activities.process_slots,
|
||||
activities.create_schedule_config,
|
||||
activities.create_slot_config,
|
||||
activities.report_schedule_orchestration,
|
||||
activities.report_slot_orchestration,
|
||||
activities.format_schedule_config,
|
||||
activities.get_last_data_timestamp,
|
||||
activities.load_latest_data,
|
||||
activities.put_last_data_timestamp,
|
||||
activities.filter_notification_alerts,
|
||||
activities.filter_notification_reports,
|
||||
activities.build_email_html,
|
||||
activities.send_email,
|
||||
activities.format_log_report,
|
||||
activities.export_data_to_postgres,
|
||||
activities.store_notification_cache,
|
||||
]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def orchestrator_worker(temporal_env, test_activities):
|
||||
"""Worker for orchestrator workflows and all activities on the default namespace."""
|
||||
with ThreadPoolExecutor(max_workers=32) as activity_executor:
|
||||
async with Worker(
|
||||
temporal_env.client,
|
||||
task_queue=ORCHESTRATOR_TASK_QUEUE,
|
||||
workflows=[
|
||||
Orchestrator,
|
||||
Alerts,
|
||||
Reports,
|
||||
LoadNotificationPackage,
|
||||
ProcessNotifications,
|
||||
],
|
||||
activities=_orchestrator_activity_list(test_activities),
|
||||
activity_executor=activity_executor,
|
||||
) as worker:
|
||||
yield worker
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def stub_workers(temporal_env):
|
||||
"""No-op workers on scouter/laborious namespaces for every managed workflow type."""
|
||||
host = temporal_host_from_env(temporal_env)
|
||||
worker_contexts: list[Worker] = []
|
||||
clients: list[Client] = []
|
||||
stub_types = list(schedule_types.keys()) + [
|
||||
'xgboost_predictions_batch',
|
||||
'xgboost_minimal_retrain',
|
||||
]
|
||||
|
||||
try:
|
||||
for namespace in MANAGED_NAMESPACES:
|
||||
client = await Client.connect(host, namespace=namespace)
|
||||
clients.append(client)
|
||||
queues = {
|
||||
build_queue_name(workflow_type, runtime)
|
||||
for workflow_type in stub_types
|
||||
for runtime in E2E_RUNTIMES
|
||||
}
|
||||
for queue in queues:
|
||||
worker = Worker(
|
||||
client,
|
||||
task_queue=queue,
|
||||
workflows=STUB_WORKFLOW_CLASSES,
|
||||
)
|
||||
await worker.__aenter__()
|
||||
worker_contexts.append(worker)
|
||||
yield worker_contexts
|
||||
finally:
|
||||
for worker in reversed(worker_contexts):
|
||||
await worker.__aexit__(None, None, None)
|
||||
33
e2e/db_schema.sql
Normal file
33
e2e/db_schema.sql
Normal file
@@ -0,0 +1,33 @@
|
||||
-- =============================================================================
|
||||
-- E2E test database schema for the ``sientia_data`` namespace.
|
||||
--
|
||||
-- SINGLE SOURCE OF TRUTH: mirrors production DDL for tables the orchestrator
|
||||
-- writes to. Any production DDL change must be pasted into this file (same
|
||||
-- pattern as sientia-dataops-laborious_temporal/e2e/db_schema.sql).
|
||||
-- =============================================================================
|
||||
|
||||
CREATE SCHEMA IF NOT EXISTS sientia_data;
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.log_report
|
||||
-- Written by ProcessNotifications via export_data_to_postgres.
|
||||
-- The table is dropped between tests so each scenario starts with a clean
|
||||
-- slate; the per-test autouse fixture re-runs this script.
|
||||
-- -----------------------------------------------------------------------------
|
||||
DROP TABLE IF EXISTS sientia_data.log_report;
|
||||
CREATE TABLE sientia_data.log_report (
|
||||
status text,
|
||||
"timestamp" timestamptz,
|
||||
groups text,
|
||||
message text,
|
||||
level text,
|
||||
notification_id text,
|
||||
block text,
|
||||
schedule text,
|
||||
pipeline text,
|
||||
project text,
|
||||
model_name text,
|
||||
model_id text,
|
||||
mail_type text,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
356
e2e/helpers.py
Normal file
356
e2e/helpers.py
Normal file
@@ -0,0 +1,356 @@
|
||||
"""
|
||||
Shared helpers for orchestrator E2E tests (Temporal workflows + Mongo + Redis + Postgres).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pymongo import MongoClient
|
||||
from redis import Redis
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Engine
|
||||
from temporalio.client import Client
|
||||
|
||||
SCENARIO_INPUTS_DIR = Path(__file__).parent / 'scenario_inputs'
|
||||
ORCHESTRATOR_TASK_QUEUE = 'orchestrator-test-queue'
|
||||
|
||||
MONGO_COLLECTIONS = (
|
||||
'notification_queue',
|
||||
'receiver_groups',
|
||||
'orchestrated_schedules',
|
||||
'pipelines',
|
||||
'opc_servers',
|
||||
'opc-servers',
|
||||
)
|
||||
|
||||
DATETIME_FORMAT_MS_WITH_TZ = '%Y-%m-%d %H:%M:%S.%f%z'
|
||||
DATETIME_FORMAT_WITH_TZ = '%Y-%m-%d %H:%M:%S%z'
|
||||
|
||||
_TIMESTAMP_MARKER_PATTERN = re.compile(r'^@now(?:([+-])(\d+)([smhd]))?$')
|
||||
_TIMESTAMP_UNIT_TO_KWARG = {'s': 'seconds', 'm': 'minutes', 'h': 'hours', 'd': 'days'}
|
||||
|
||||
|
||||
def _resolve_timestamp_marker(value: Any) -> Any:
|
||||
"""
|
||||
Convert ``@now`` / ``@now-1h`` markers into timezone-aware datetimes.
|
||||
|
||||
Args:
|
||||
value: Any JSON value. Only strings matching the marker pattern are converted.
|
||||
|
||||
Return:
|
||||
Any: The resolved datetime or the original value unchanged.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
match = _TIMESTAMP_MARKER_PATTERN.match(value)
|
||||
if not match:
|
||||
return value
|
||||
sign, amount, unit = match.groups()
|
||||
now = datetime.now(UTC)
|
||||
if sign is None:
|
||||
return now
|
||||
delta = timedelta(**{_TIMESTAMP_UNIT_TO_KWARG[unit]: int(amount)})
|
||||
return now + delta if sign == '+' else now - delta
|
||||
|
||||
|
||||
def _resolve_payload(payload: Any) -> Any:
|
||||
"""Recursively walk a JSON-like structure resolving ``@now`` timestamp markers."""
|
||||
if isinstance(payload, dict):
|
||||
return {key: _resolve_payload(value) for key, value in payload.items()}
|
||||
if isinstance(payload, list):
|
||||
return [_resolve_payload(item) for item in payload]
|
||||
return _resolve_timestamp_marker(payload)
|
||||
|
||||
|
||||
def load_scenario_input(name: str, **overrides: Any) -> dict[str, Any]:
|
||||
"""
|
||||
Load a scenario JSON file from e2e/scenario_inputs and apply overrides.
|
||||
|
||||
Strings matching ``@now`` or ``@now[+-]<int>[smhd]`` (anywhere in the payload)
|
||||
are converted to timezone-aware ``datetime`` instances. This lets scenario
|
||||
files declare relative timestamps such as ``"updated_at": "@now-1h"``.
|
||||
|
||||
Args:
|
||||
name: File name (with or without .json suffix).
|
||||
**overrides: Top-level keys to replace in the loaded dict.
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Scenario payload with timestamp markers resolved.
|
||||
"""
|
||||
file_name = name if name.endswith('.json') else f'{name}.json'
|
||||
file_path = SCENARIO_INPUTS_DIR / file_name
|
||||
with file_path.open('r', encoding='utf-8') as handle:
|
||||
payload = json.load(handle)
|
||||
payload = _resolve_payload(payload)
|
||||
payload.update(overrides)
|
||||
return payload
|
||||
|
||||
|
||||
def make_workflow_id(prefix: str) -> str:
|
||||
"""Build a unique workflow id using a prefix and UUID suffix."""
|
||||
return f'{prefix}-{uuid.uuid4().hex[:12]}'
|
||||
|
||||
|
||||
async def start_and_await_workflow(
|
||||
client: Client,
|
||||
workflow_run,
|
||||
input_data: dict[str, Any],
|
||||
workflow_id: str,
|
||||
*,
|
||||
task_queue: str = ORCHESTRATOR_TASK_QUEUE,
|
||||
timeout: float = 120.0,
|
||||
) -> Any:
|
||||
"""
|
||||
Start a workflow and wait for its result.
|
||||
|
||||
Args:
|
||||
client: Temporal client (default namespace).
|
||||
workflow_run: Workflow run method (e.g. Orchestrator.run).
|
||||
input_data: Workflow input payload.
|
||||
workflow_id: Unique workflow id.
|
||||
task_queue: Task queue for the orchestrator worker.
|
||||
timeout: Max seconds to wait for completion.
|
||||
|
||||
Return:
|
||||
Workflow result value.
|
||||
"""
|
||||
handle = await client.start_workflow(
|
||||
workflow_run,
|
||||
input_data,
|
||||
id=workflow_id,
|
||||
task_queue=task_queue,
|
||||
)
|
||||
return await asyncio.wait_for(handle.result(), timeout=timeout)
|
||||
|
||||
|
||||
def seed_pipelines(
|
||||
mongo_uri: str,
|
||||
database: str,
|
||||
pipelines: list[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Insert pipeline documents into the test Mongo database."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
collection = client[database]['pipelines']
|
||||
if pipelines:
|
||||
collection.insert_many(pipelines)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def seed_opc_servers(
|
||||
mongo_uri: str,
|
||||
database: str,
|
||||
servers: list[dict[str, Any]],
|
||||
*,
|
||||
collection: str = 'opc_servers',
|
||||
) -> None:
|
||||
"""Insert OPC server documents into the test Mongo database."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
coll = client[database][collection]
|
||||
if servers:
|
||||
coll.insert_many(servers)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def seed_receiver_groups(
|
||||
mongo_uri: str,
|
||||
database: str,
|
||||
groups: list[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Insert receiver group documents into the test Mongo database."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
collection = client[database]['receiver_groups']
|
||||
if groups:
|
||||
collection.insert_many(groups)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def seed_notifications(
|
||||
mongo_uri: str,
|
||||
database: str,
|
||||
notifications: list[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Insert notification_queue documents into the test Mongo database."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
collection = client[database]['notification_queue']
|
||||
if notifications:
|
||||
collection.insert_many(notifications)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def seed_orchestrated_schedules(
|
||||
mongo_uri: str,
|
||||
database: str,
|
||||
schedules: list[dict[str, Any]],
|
||||
) -> None:
|
||||
"""Insert orchestrated_schedules tracking documents."""
|
||||
client = MongoClient(mongo_uri)
|
||||
try:
|
||||
collection = client[database]['orchestrated_schedules']
|
||||
if schedules:
|
||||
collection.insert_many(schedules)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def seed_opc_slots(redis_client: Redis, slots: dict[str, str]) -> None:
|
||||
"""Write OPC slot keys (slot:opc_tags:*) in Redis."""
|
||||
for key, value in slots.items():
|
||||
redis_client.set(key, value)
|
||||
|
||||
|
||||
def seed_active_ingestors(redis_client: Redis, ingestor_keys: list[str]) -> None:
|
||||
"""Seed heartbeat:ingestor:* keys so load_active_ingestors returns ingestors."""
|
||||
for key in ingestor_keys:
|
||||
redis_client.set(key, '1')
|
||||
|
||||
|
||||
def seed_last_timestamp(redis_client: Redis, mail_type: str, value: str) -> None:
|
||||
"""
|
||||
Set notification_last_timestamp for a mail type, JSON-encoded.
|
||||
|
||||
Values must be JSON-encoded so ``redis_repository.get`` (which calls
|
||||
``json.loads`` on the raw payload) can deserialize them. The value
|
||||
must follow the exact format ``sientia_do.notifications.models.Notification``
|
||||
writes into ``notification_queue.timestamp``: ``DATETIME_FORMAT_WITH_TZ``
|
||||
(e.g. ``"2026-05-22 16:47:02+0000"``) — no microseconds and no colon in
|
||||
the timezone offset.
|
||||
|
||||
Args:
|
||||
redis_client: Redis client connected to the test instance.
|
||||
mail_type: Mail type identifier (e.g. ``"Alerts"``, ``"Reports"``).
|
||||
value: Timestamp string in ``DATETIME_FORMAT_WITH_TZ``
|
||||
(e.g. ``"2024-06-01 10:30:00+0000"``).
|
||||
"""
|
||||
redis_client.set(f'notification_last_timestamp:{mail_type}', json.dumps(value))
|
||||
|
||||
|
||||
def seed_notification_cache(
|
||||
redis_client: Redis,
|
||||
trigger: str,
|
||||
notification_id: str,
|
||||
*,
|
||||
sent_at: str | None = None,
|
||||
ttl: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Pre-seed alerts sent cache entry, JSON-encoded.
|
||||
|
||||
The value must be JSON-encoded because ``filter_notification_alerts``
|
||||
reads via ``redis_repository.get`` (which applies ``json.loads``) and
|
||||
parses the resulting string with ``DATETIME_FORMAT_MS_WITH_TZ``.
|
||||
|
||||
Args:
|
||||
redis_client: Redis client connected to the test instance.
|
||||
trigger: Schedule/trigger name used to compose the cache key.
|
||||
notification_id: Notification id used to compose the cache key.
|
||||
sent_at: Optional timestamp string in ``DATETIME_FORMAT_MS_WITH_TZ``.
|
||||
ttl: Optional TTL in seconds for the cache entry.
|
||||
"""
|
||||
key = f'{trigger}:{notification_id}'
|
||||
value = sent_at or datetime.now(UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ)
|
||||
encoded = json.dumps(value)
|
||||
if ttl is not None:
|
||||
redis_client.set(key, encoded, ex=ttl)
|
||||
else:
|
||||
redis_client.set(key, encoded)
|
||||
|
||||
|
||||
def count_log_report_rows(engine: Engine, mail_type: str | None = None) -> int:
|
||||
"""Count rows in sientia_data.log_report, optionally filtered by mail_type."""
|
||||
sql = 'SELECT COUNT(*) FROM sientia_data.log_report'
|
||||
params: dict[str, Any] = {}
|
||||
if mail_type is not None:
|
||||
sql += ' WHERE mail_type = :mail_type'
|
||||
params['mail_type'] = mail_type
|
||||
with engine.connect() as conn:
|
||||
return int(conn.execute(text(sql), params).scalar() or 0)
|
||||
|
||||
|
||||
def fetch_log_report(engine: Engine, mail_type: str | None = None) -> list[dict[str, Any]]:
|
||||
"""Fetch log_report rows as dicts."""
|
||||
sql = 'SELECT * FROM sientia_data.log_report'
|
||||
params: dict[str, Any] = {}
|
||||
if mail_type is not None:
|
||||
sql += ' WHERE mail_type = :mail_type'
|
||||
params['mail_type'] = mail_type
|
||||
with engine.connect() as conn:
|
||||
rows = conn.execute(text(sql), params).mappings().all()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
def default_notification(
|
||||
*,
|
||||
notification_id: str,
|
||||
level: str = 'ERROR',
|
||||
timestamp: str | None = None,
|
||||
model_name: str = 'model-a',
|
||||
model_id: str = '1',
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Build a minimal notification_queue document mirroring production layout.
|
||||
|
||||
The ``timestamp`` field is stored as a string in ``DATETIME_FORMAT_WITH_TZ``
|
||||
because that is exactly what ``sientia_do.notifications.models.Notification``
|
||||
writes into ``notification_queue`` in production (``now().strftime(
|
||||
DATETIME_FORMAT_WITH_TZ)``). Tests intentionally use this same format so we
|
||||
surface, rather than hide, real production behavior in downstream
|
||||
activities.
|
||||
|
||||
Args:
|
||||
notification_id: Unique identifier for the notification.
|
||||
level: Notification level (e.g. ``"ERROR"``, ``"WARNING"``).
|
||||
timestamp: Optional production-format timestamp string. ``None`` falls
|
||||
back to a fixed sample value.
|
||||
model_name: Model name attached to the notification.
|
||||
model_id: Model id attached to the notification.
|
||||
|
||||
Return:
|
||||
dict[str, Any]: A notification document ready for insertion.
|
||||
"""
|
||||
ts = timestamp if timestamp is not None else datetime(
|
||||
2024, 6, 1, 12, 0, 0, tzinfo=UTC
|
||||
).strftime(DATETIME_FORMAT_WITH_TZ)
|
||||
return {
|
||||
'notification_id': notification_id,
|
||||
'level': level,
|
||||
'timestamp': ts,
|
||||
'message': f'{level} on {model_name}',
|
||||
'trigger': 'test-schedule',
|
||||
'block': 'test-block',
|
||||
'pipeline': 'test-pipeline',
|
||||
'project': 'orchestrator-e2e',
|
||||
'model_name': model_name,
|
||||
'model_id': model_id,
|
||||
}
|
||||
|
||||
|
||||
def default_receiver_group(
|
||||
*,
|
||||
group_name: str = 'admins',
|
||||
members: list[str] | None = None,
|
||||
levels: list[str] | None = None,
|
||||
contents: list[str] | None = None,
|
||||
ignore_models: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Minimal active receiver_groups document."""
|
||||
return {
|
||||
'group_name': group_name,
|
||||
'active': True,
|
||||
'members': members or ['admin@example.com'],
|
||||
'levels': levels or ['ERROR', 'WARNING', 'INFO'],
|
||||
'contents': contents or ['core_alerts', 'persistent_alerts', 'reports'],
|
||||
'ignore_models': ignore_models or [],
|
||||
}
|
||||
5
e2e/scenario_inputs/alerts_duplicate.json
Normal file
5
e2e/scenario_inputs/alerts_duplicate.json
Normal file
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"schedule_name": "alerts-e2e-dup",
|
||||
"notification_ttl": 300,
|
||||
"sent_ttl": 600
|
||||
}
|
||||
5
e2e/scenario_inputs/alerts_empty.json
Normal file
5
e2e/scenario_inputs/alerts_empty.json
Normal file
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"schedule_name": "alerts-e2e-empty",
|
||||
"notification_ttl": 300,
|
||||
"sent_ttl": 600
|
||||
}
|
||||
5
e2e/scenario_inputs/alerts_happy_path.json
Normal file
5
e2e/scenario_inputs/alerts_happy_path.json
Normal file
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"schedule_name": "alerts-e2e",
|
||||
"notification_ttl": 300,
|
||||
"sent_ttl": 600
|
||||
}
|
||||
5
e2e/scenario_inputs/alerts_persistent.json
Normal file
5
e2e/scenario_inputs/alerts_persistent.json
Normal file
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"schedule_name": "alerts-e2e-persistent",
|
||||
"notification_ttl": 1,
|
||||
"sent_ttl": 600
|
||||
}
|
||||
34
e2e/scenario_inputs/orchestrator_conflict.json
Normal file
34
e2e/scenario_inputs/orchestrator_conflict.json
Normal file
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-conflict",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "conflict-pred",
|
||||
"workflow_type": "predictions_batch",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"query": "SELECT 1",
|
||||
"write_tags": [
|
||||
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"]
|
||||
}
|
||||
34
e2e/scenario_inputs/orchestrator_create_only.json
Normal file
34
e2e/scenario_inputs/orchestrator_create_only.json
Normal file
@@ -0,0 +1,34 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-create",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "create-only-pred",
|
||||
"workflow_type": "predictions_batch",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"query": "SELECT 1",
|
||||
"write_tags": [
|
||||
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"]
|
||||
}
|
||||
31
e2e/scenario_inputs/orchestrator_delete_only.json
Normal file
31
e2e/scenario_inputs/orchestrator_delete_only.json
Normal file
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-delete",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "delete-me",
|
||||
"workflow_type": "drift",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now-1h",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"interval_minutes": 60
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"]
|
||||
}
|
||||
16
e2e/scenario_inputs/orchestrator_empty.json
Normal file
16
e2e/scenario_inputs/orchestrator_empty.json
Normal file
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-empty",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [],
|
||||
"opc_servers": [],
|
||||
"active_ingestors": []
|
||||
}
|
||||
46
e2e/scenario_inputs/orchestrator_happy_path.json
Normal file
46
e2e/scenario_inputs/orchestrator_happy_path.json
Normal file
@@ -0,0 +1,46 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "pred-legacy",
|
||||
"workflow_type": "predictions_batch",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"query": "SELECT 1",
|
||||
"write_tags": [
|
||||
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"schedule_name": "drift-gpu",
|
||||
"workflow_type": "drift",
|
||||
"runtime": "gpu",
|
||||
"model_id": "model-2",
|
||||
"model": {"name": "Model model-2"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"interval_minutes": 60
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true, "name": "opc-1"}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"]
|
||||
}
|
||||
41
e2e/scenario_inputs/orchestrator_noop.json
Normal file
41
e2e/scenario_inputs/orchestrator_noop.json
Normal file
@@ -0,0 +1,41 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-noop",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "noop-pred",
|
||||
"workflow_type": "predictions_batch",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now-1h",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"query": "SELECT 1",
|
||||
"write_tags": [
|
||||
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"],
|
||||
"orchestrated_schedules": [
|
||||
{
|
||||
"schedule_name": "noop-pred",
|
||||
"namespace": "laborious",
|
||||
"updated_at": "@now-1h"
|
||||
}
|
||||
]
|
||||
}
|
||||
45
e2e/scenario_inputs/orchestrator_ttl_index.json
Normal file
45
e2e/scenario_inputs/orchestrator_ttl_index.json
Normal file
@@ -0,0 +1,45 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-ttl",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "scouter-ttl",
|
||||
"workflow_type": "scouter",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"read_tags": [
|
||||
{
|
||||
"server_id": "srv-1",
|
||||
"tag_name": "Read1",
|
||||
"tag_address": "ns=2;s=Read1",
|
||||
"aggr_func": "lts",
|
||||
"frequency": 1000
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{
|
||||
"id": "srv-1",
|
||||
"active": true,
|
||||
"server_name": "opc-1",
|
||||
"url": "opc.tcp://localhost:4840",
|
||||
"uri": "urn:opcfoundation:UA:DemoServer"
|
||||
}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"]
|
||||
}
|
||||
41
e2e/scenario_inputs/orchestrator_update_only.json
Normal file
41
e2e/scenario_inputs/orchestrator_update_only.json
Normal file
@@ -0,0 +1,41 @@
|
||||
{
|
||||
"workflow_input": {
|
||||
"schedule_name": "orchestrator-e2e-update",
|
||||
"pipelines_query": {
|
||||
"collection": "pipelines",
|
||||
"aggregation": [{"$match": {"active": true}}]
|
||||
},
|
||||
"opc_servers_query": {
|
||||
"collection": "opc_servers",
|
||||
"filters": {"active": true}
|
||||
}
|
||||
},
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "update-pred",
|
||||
"workflow_type": "predictions_batch",
|
||||
"runtime": "legacy",
|
||||
"model_id": "model-1",
|
||||
"model": {"name": "Model model-1"},
|
||||
"active": true,
|
||||
"updated_at": "@now",
|
||||
"frequency": "5m",
|
||||
"offset": "0m",
|
||||
"query": "SELECT 1",
|
||||
"write_tags": [
|
||||
{"server_id": "srv-1", "type": "prediction", "addr": "ns=2;s=Tag1"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"opc_servers": [
|
||||
{"id": "srv-1", "active": true}
|
||||
],
|
||||
"active_ingestors": ["heartbeat:ingestor:1"],
|
||||
"orchestrated_schedules": [
|
||||
{
|
||||
"schedule_name": "update-pred",
|
||||
"namespace": "laborious",
|
||||
"updated_at": "@now-1h"
|
||||
}
|
||||
]
|
||||
}
|
||||
3
e2e/scenario_inputs/reports_empty.json
Normal file
3
e2e/scenario_inputs/reports_empty.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"schedule_name": "reports-e2e-empty"
|
||||
}
|
||||
3
e2e/scenario_inputs/reports_happy_path.json
Normal file
3
e2e/scenario_inputs/reports_happy_path.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"schedule_name": "reports-e2e"
|
||||
}
|
||||
3
e2e/scenario_inputs/reports_multi_level.json
Normal file
3
e2e/scenario_inputs/reports_multi_level.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"schedule_name": "reports-e2e-levels"
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"schedule_name": "load-pkg-e2e",
|
||||
"mail_type": "Alerts",
|
||||
"base_data_filter": {"level": "ERROR"}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"schedule_name": "process-notif-e2e",
|
||||
"mail_type": "Alerts",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "log_report"
|
||||
}
|
||||
111
e2e/scenarios.md
Normal file
111
e2e/scenarios.md
Normal file
@@ -0,0 +1,111 @@
|
||||
# E2E Scenario Documentation — Orchestrator
|
||||
|
||||
Functional reference for orchestrator E2E scenarios. Tests live under `e2e/`, use `@pytest.mark.e2e`, and run with:
|
||||
|
||||
```bash
|
||||
pytest e2e/ --override-ini testpaths=e2e -m e2e
|
||||
```
|
||||
|
||||
## Execution context
|
||||
|
||||
- MongoDB, Redis, PostgreSQL: testcontainers (session-scoped).
|
||||
- SMTP: in-process `aiosmtpd` (`e2e/smtp_test_server.py`).
|
||||
- Temporal: `WorkflowEnvironment.start_local()` with stub workers on `scouter` / `laborious`.
|
||||
- Production code under `orchestrator/**` is not mocked; only `Logger` may be a `MagicMock`.
|
||||
|
||||
---
|
||||
|
||||
## 1. Orchestrator workflow
|
||||
|
||||
Source: `e2e/test_orchestrator_main_workflow.py`
|
||||
|
||||
### 1.1.1 Happy path
|
||||
|
||||
Pipelines in Mongo → schedules created in correct namespace/task queue, Redis slots written, `orchestrated_schedules` updated.
|
||||
|
||||
### 1.2.1 No-op tick
|
||||
|
||||
Mongo, Redis, and Temporal already match desired state → no new schedules or slot writes.
|
||||
|
||||
### 1.3.1 Create-only
|
||||
|
||||
New pipeline only → schedules created, timestamps inserted.
|
||||
|
||||
### 1.3.2 Update-only
|
||||
|
||||
Existing pipeline with newer `updated_at` → schedule updated in Temporal.
|
||||
|
||||
### 1.3.3 Delete-only
|
||||
|
||||
Pipeline removed from Mongo → schedule deleted from Temporal.
|
||||
|
||||
### 1.4.1 Conflict ordering
|
||||
|
||||
Pipeline update and slot delete on same OPC server → slot insert before delete (production ordering).
|
||||
|
||||
### 1.5.1 Empty pipelines
|
||||
|
||||
No active pipelines → orphan schedules removed, no new orchestration writes.
|
||||
|
||||
### 1.6.1 TTL index bootstrap
|
||||
|
||||
First run creates TTL index on notification collection used by scouter pipelines.
|
||||
|
||||
---
|
||||
|
||||
## 2. Alerts workflow
|
||||
|
||||
Source: `e2e/test_alerts_main_workflow.py`
|
||||
|
||||
### A.1.1 Happy path
|
||||
|
||||
ERROR notification → one SMTP message, one `log_report` row, Redis cache key.
|
||||
|
||||
### A.1.2 TTL duplicate suppression
|
||||
|
||||
Second run with same data and cache seeded → no extra email or log row.
|
||||
|
||||
### A.1.3 Persistent escalation
|
||||
|
||||
Alert past `notification_ttl` with cache cleared → new email sent.
|
||||
|
||||
### A.2.1 Group filtering
|
||||
|
||||
Receiver group `levels` / `ignore_models` honored.
|
||||
|
||||
### A.3.1 Empty queue
|
||||
|
||||
No notifications → no SMTP, no Postgres row.
|
||||
|
||||
---
|
||||
|
||||
## 3. Reports workflow
|
||||
|
||||
Source: `e2e/test_reports_main_workflow.py`
|
||||
|
||||
### R.1.1 Happy path
|
||||
|
||||
Mixed ERROR/WARNING/INFO → one HTML email with all section headings.
|
||||
|
||||
### R.1.2 Per-level rendering
|
||||
|
||||
Single-level notifications → only matching section in HTML body.
|
||||
|
||||
### R.2.1 Empty queue
|
||||
|
||||
No notifications → no SMTP, no Postgres row.
|
||||
|
||||
---
|
||||
|
||||
## 4. Subworkflows
|
||||
|
||||
### LoadNotificationPackage — `e2e/test_subworkflow_load_notification_package.py`
|
||||
|
||||
- No prior Redis timestamp → all matching notifications returned, max timestamp stored.
|
||||
- Prior timestamp → only newer notifications returned.
|
||||
- Empty Mongo → no Redis timestamp write.
|
||||
|
||||
### ProcessNotifications — `e2e/test_subworkflow_process_notifications.py`
|
||||
|
||||
- Full round trip: HTML → SMTP → `log_report` in Postgres.
|
||||
- Empty receiver groups → `{}`.
|
||||
110
e2e/smtp_test_server.py
Normal file
110
e2e/smtp_test_server.py
Normal file
@@ -0,0 +1,110 @@
|
||||
"""In-process SMTP server for E2E email assertions."""
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
import threading
|
||||
import time
|
||||
from email import message_from_bytes, policy
|
||||
from email.message import EmailMessage, Message
|
||||
|
||||
from aiosmtpd.controller import Controller
|
||||
|
||||
|
||||
class _CaptureHandler:
|
||||
"""
|
||||
aiosmtpd handler that parses every incoming message into an ``EmailMessage``.
|
||||
|
||||
The default policy yields a ``Message`` instance, which loses structure
|
||||
when wrapped into a new ``EmailMessage``. Parsing with ``policy.default``
|
||||
keeps multipart payloads intact so tests can introspect the HTML body
|
||||
via ``walk()``/``get_payload(decode=True)``.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.messages: list[EmailMessage | Message] = []
|
||||
|
||||
async def handle_DATA(self, server, session, envelope):
|
||||
message = message_from_bytes(envelope.content, policy=policy.default)
|
||||
self.messages.append(message)
|
||||
return '250 OK'
|
||||
|
||||
|
||||
def _reserve_port(host: str = '127.0.0.1') -> int:
|
||||
"""Reserve a free TCP port on the given host."""
|
||||
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
probe.bind((host, 0))
|
||||
port = probe.getsockname()[1]
|
||||
probe.close()
|
||||
return port
|
||||
|
||||
|
||||
class SmtpTestServer:
|
||||
"""
|
||||
Wraps aiosmtpd Controller with a dedicated asyncio loop thread for pytest compatibility.
|
||||
|
||||
Attributes:
|
||||
host: Bind host (127.0.0.1).
|
||||
port: Listening port after start().
|
||||
messages: Captured outbound messages.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.host = '127.0.0.1'
|
||||
self.port: int | None = None
|
||||
self._handler = _CaptureHandler()
|
||||
self.messages: list[EmailMessage | Message] = self._handler.messages
|
||||
self._controller: Controller | None = None
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the SMTP controller on a reserved port in a background event loop."""
|
||||
self.port = _reserve_port(self.host)
|
||||
self._loop = asyncio.new_event_loop()
|
||||
self._controller = Controller(
|
||||
self._handler,
|
||||
hostname=self.host,
|
||||
port=self.port,
|
||||
loop=self._loop,
|
||||
ready_timeout=30,
|
||||
)
|
||||
|
||||
def _run():
|
||||
asyncio.set_event_loop(self._loop)
|
||||
if self._controller is None:
|
||||
raise RuntimeError('SMTP controller not initialized')
|
||||
self._controller.start()
|
||||
|
||||
self._thread = threading.Thread(target=_run, name='e2e-smtp', daemon=True)
|
||||
self._thread.start()
|
||||
self._wait_until_ready()
|
||||
|
||||
def _wait_until_ready(self, timeout: float = 10.0) -> None:
|
||||
"""Poll until the SMTP listener accepts TCP connections."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if self.port is None:
|
||||
time.sleep(0.05)
|
||||
continue
|
||||
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
try:
|
||||
if probe.connect_ex((self.host, self.port)) == 0:
|
||||
return
|
||||
finally:
|
||||
probe.close()
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError('SMTP test server did not become ready')
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the controller and background event loop."""
|
||||
if self._controller is not None:
|
||||
self._controller.stop()
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=5)
|
||||
self._controller = None
|
||||
self._loop = None
|
||||
self._thread = None
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Remove all captured messages."""
|
||||
self.messages.clear()
|
||||
73
e2e/stub_workflows.py
Normal file
73
e2e/stub_workflows.py
Normal file
@@ -0,0 +1,73 @@
|
||||
"""No-op Temporal workflows for managed scouter/laborious namespaces in E2E."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
|
||||
@workflow.defn(name='scouter')
|
||||
class ScouterStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='pi_web_api_scouter')
|
||||
class PiWebApiScouterStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='predictions_batch')
|
||||
class PredictionsBatchStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='xgboost_predictions_batch')
|
||||
class XgboostPredictionsBatchStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='drift')
|
||||
class DriftStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='simple_metrics')
|
||||
class SimpleMetricsStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='minimal_retrain')
|
||||
class MinimalRetrainStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@workflow.defn(name='xgboost_minimal_retrain')
|
||||
class XgboostMinimalRetrainStub:
|
||||
@workflow.run
|
||||
async def run(self, _input_data: dict[str, Any]) -> None:
|
||||
return None
|
||||
|
||||
|
||||
STUB_WORKFLOW_CLASSES = [
|
||||
ScouterStub,
|
||||
PiWebApiScouterStub,
|
||||
PredictionsBatchStub,
|
||||
XgboostPredictionsBatchStub,
|
||||
DriftStub,
|
||||
SimpleMetricsStub,
|
||||
MinimalRetrainStub,
|
||||
XgboostMinimalRetrainStub,
|
||||
]
|
||||
116
e2e/test_alerts_main_workflow.py
Normal file
116
e2e/test_alerts_main_workflow.py
Normal file
@@ -0,0 +1,116 @@
|
||||
"""E2E tests for the Alerts main workflow."""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
from redis import Redis
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
|
||||
from e2e.conftest import E2E_DATABASE
|
||||
from e2e.helpers import (
|
||||
DATETIME_FORMAT_MS_WITH_TZ,
|
||||
count_log_report_rows,
|
||||
default_notification,
|
||||
default_receiver_group,
|
||||
fetch_log_report,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
seed_notification_cache,
|
||||
seed_notifications,
|
||||
seed_receiver_groups,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_a_1_1_happy_path(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""A.1.1: ERROR alert sends email, writes log_report, caches notification."""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
seed_notifications(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[default_notification(notification_id='alert-1', level='ERROR')],
|
||||
)
|
||||
|
||||
input_data = load_scenario_input('alerts_happy_path.json')
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Alerts.run,
|
||||
input_data,
|
||||
make_workflow_id('alerts-happy'),
|
||||
)
|
||||
|
||||
assert len(smtp_server.messages) == 1
|
||||
assert count_log_report_rows(postgres_engine, 'Alerts') == 1
|
||||
rows = fetch_log_report(postgres_engine, 'Alerts')
|
||||
assert rows[0]['mail_type'] == 'Alerts'
|
||||
assert redis_client.get('test-schedule:alert-1') is not None or redis_client.keys('*alert-1*')
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_a_1_2_duplicate_suppressed(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""A.1.2: Cached notification is not emailed twice within sent_ttl."""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
notif = default_notification(notification_id='dup-1', level='ERROR')
|
||||
seed_notifications(mongo_uri, E2E_DATABASE, [notif])
|
||||
seed_notification_cache(
|
||||
redis_client,
|
||||
notif['trigger'],
|
||||
notif['notification_id'],
|
||||
sent_at=datetime.now(UTC).strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||
ttl=600,
|
||||
)
|
||||
|
||||
input_data = load_scenario_input('alerts_duplicate.json')
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Alerts.run,
|
||||
input_data,
|
||||
make_workflow_id('alerts-dup'),
|
||||
)
|
||||
|
||||
assert len(smtp_server.messages) == 0
|
||||
assert count_log_report_rows(postgres_engine, 'Alerts') == 0
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_a_3_1_empty_queue(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
):
|
||||
"""A.3.1: Empty notification queue short-circuits."""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Alerts.run,
|
||||
load_scenario_input('alerts_empty.json'),
|
||||
make_workflow_id('alerts-empty'),
|
||||
)
|
||||
|
||||
assert len(smtp_server.messages) == 0
|
||||
assert count_log_report_rows(postgres_engine, 'Alerts') == 0
|
||||
226
e2e/test_orchestrator_main_workflow.py
Normal file
226
e2e/test_orchestrator_main_workflow.py
Normal file
@@ -0,0 +1,226 @@
|
||||
"""E2E tests for the Orchestrator main workflow."""
|
||||
|
||||
import pytest
|
||||
from pymongo import MongoClient
|
||||
from redis import Redis
|
||||
from sientia_do.temporal.worker.prepare_worker import build_queue_name
|
||||
from temporalio.client import Client, Schedule, ScheduleActionStartWorkflow, ScheduleSpec
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
|
||||
from e2e.conftest import E2E_DATABASE
|
||||
from e2e.helpers import (
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
seed_active_ingestors,
|
||||
seed_opc_servers,
|
||||
seed_orchestrated_schedules,
|
||||
seed_pipelines,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
|
||||
|
||||
async def _schedule_ids(host: str, namespace: str) -> list[str]:
|
||||
"""Return the list of Temporal schedule ids in a given namespace."""
|
||||
client = await Client.connect(host, namespace=namespace)
|
||||
return [schedule.id async for schedule in await client.list_schedules()]
|
||||
|
||||
|
||||
def _apply_orchestrator_seeds(
|
||||
scenario: dict,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
) -> None:
|
||||
"""
|
||||
Seed Mongo and Redis with the pipelines/opc_servers/ingestors declared in a scenario.
|
||||
|
||||
Args:
|
||||
scenario: Scenario payload returned by ``load_scenario_input``.
|
||||
mongo_uri: Mongo connection string for the test database.
|
||||
redis_client: Redis client connected to the test instance.
|
||||
"""
|
||||
seed_pipelines(mongo_uri, E2E_DATABASE, scenario.get('pipelines', []))
|
||||
seed_opc_servers(mongo_uri, E2E_DATABASE, scenario.get('opc_servers', []))
|
||||
seed_active_ingestors(redis_client, scenario.get('active_ingestors', []))
|
||||
if scenario.get('orchestrated_schedules'):
|
||||
seed_orchestrated_schedules(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
scenario['orchestrated_schedules'],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_1_1_1_happy_path(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
temporal_host: str,
|
||||
):
|
||||
"""1.1.1: Creates schedules, slots, and orchestrated_schedules entries."""
|
||||
scenario = load_scenario_input('orchestrator_happy_path')
|
||||
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-happy'),
|
||||
)
|
||||
|
||||
mongo = MongoClient(mongo_uri)
|
||||
try:
|
||||
tracked = list(mongo[E2E_DATABASE]['orchestrated_schedules'].find())
|
||||
names = {doc['schedule_name'] for doc in tracked}
|
||||
assert names, f'Expected orchestrated_schedules rows, got {tracked}'
|
||||
finally:
|
||||
mongo.close()
|
||||
|
||||
laborious_schedules = await _schedule_ids(temporal_host, 'laborious')
|
||||
assert 'pred-legacy' in laborious_schedules or 'drift-gpu' in laborious_schedules, (
|
||||
f'Expected Temporal schedules in laborious, got {laborious_schedules}'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_1_3_1_create_only(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
temporal_host: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""1.3.1: New pipeline creates a Temporal schedule."""
|
||||
scenario = load_scenario_input('orchestrator_create_only')
|
||||
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-create'),
|
||||
)
|
||||
|
||||
assert 'create-only-pred' in await _schedule_ids(temporal_host, 'laborious')
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_1_3_3_delete_only(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
temporal_host: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""1.3.3: Removing pipeline deletes Temporal schedule."""
|
||||
scenario = load_scenario_input('orchestrator_delete_only')
|
||||
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-seed-delete'),
|
||||
)
|
||||
|
||||
client = await Client.connect(temporal_host, namespace='laborious')
|
||||
try:
|
||||
handle = client.get_schedule_handle('delete-me')
|
||||
await handle.describe()
|
||||
schedule_exists = True
|
||||
except Exception:
|
||||
schedule_exists = False
|
||||
|
||||
assert schedule_exists
|
||||
|
||||
mongo = MongoClient(mongo_uri)
|
||||
try:
|
||||
mongo[E2E_DATABASE]['pipelines'].delete_many({})
|
||||
finally:
|
||||
mongo.close()
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-delete'),
|
||||
)
|
||||
|
||||
assert 'delete-me' not in await _schedule_ids(temporal_host, 'laborious')
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_1_5_1_empty_pipelines(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
temporal_host: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""1.5.1: No pipelines → no orchestrated_schedules documents."""
|
||||
scenario = load_scenario_input('orchestrator_empty')
|
||||
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
|
||||
|
||||
client = await Client.connect(temporal_host, namespace='laborious')
|
||||
await client.create_schedule(
|
||||
'orphan-schedule',
|
||||
Schedule(
|
||||
action=ScheduleActionStartWorkflow(
|
||||
'drift',
|
||||
{},
|
||||
id='orphan-schedule-run',
|
||||
task_queue=build_queue_name('drift', 'legacy'),
|
||||
),
|
||||
spec=ScheduleSpec(),
|
||||
),
|
||||
)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-empty'),
|
||||
)
|
||||
|
||||
mongo = MongoClient(mongo_uri)
|
||||
try:
|
||||
assert mongo[E2E_DATABASE]['orchestrated_schedules'].count_documents({}) == 0
|
||||
finally:
|
||||
mongo.close()
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_1_6_1_ttl_index_bootstrap(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""1.6.1: Scouter pipeline triggers TTL index on notification_queue."""
|
||||
scenario = load_scenario_input('orchestrator_ttl_index')
|
||||
_apply_orchestrator_seeds(scenario, mongo_uri, redis_client)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Orchestrator.run,
|
||||
scenario['workflow_input'],
|
||||
make_workflow_id('orchestrator-ttl'),
|
||||
)
|
||||
|
||||
mongo = MongoClient(mongo_uri)
|
||||
try:
|
||||
indexes = mongo[E2E_DATABASE]['raw_scouter-ttl'].index_information()
|
||||
assert any('expireAfterSeconds' in info for info in indexes.values())
|
||||
finally:
|
||||
mongo.close()
|
||||
167
e2e/test_reports_main_workflow.py
Normal file
167
e2e/test_reports_main_workflow.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""E2E tests for the Reports main workflow."""
|
||||
|
||||
from email.message import EmailMessage, Message
|
||||
|
||||
import pytest
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
|
||||
from e2e.conftest import E2E_DATABASE
|
||||
from e2e.helpers import (
|
||||
count_log_report_rows,
|
||||
default_notification,
|
||||
default_receiver_group,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
seed_notifications,
|
||||
seed_receiver_groups,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from orchestrator.workflows.reports import Reports
|
||||
|
||||
|
||||
def _extract_html_body(message: EmailMessage | Message) -> str:
|
||||
"""
|
||||
Return the text/html portion of an email message as a decoded string.
|
||||
|
||||
Walks every part looking for the first text/html payload, decoding it
|
||||
according to the part's transfer encoding and charset. Falls back to
|
||||
the message's own ``get_content``/raw payload when no HTML part is
|
||||
present so callers can still inspect plain-text reports.
|
||||
|
||||
Args:
|
||||
message: Captured email message returned by the test SMTP server.
|
||||
|
||||
Return:
|
||||
str: HTML body content, or an empty string when nothing decodable
|
||||
is found.
|
||||
"""
|
||||
if message.is_multipart():
|
||||
for part in message.walk():
|
||||
if part.get_content_type() != 'text/html':
|
||||
continue
|
||||
payload = part.get_payload(decode=True)
|
||||
if payload is None:
|
||||
continue
|
||||
charset = part.get_content_charset() or 'utf-8'
|
||||
return payload.decode(charset, errors='replace')
|
||||
payload = message.get_payload(decode=True)
|
||||
if payload is not None:
|
||||
charset = message.get_content_charset() or 'utf-8'
|
||||
return payload.decode(charset, errors='replace')
|
||||
try:
|
||||
return message.get_content()
|
||||
except (AttributeError, KeyError):
|
||||
return str(message.get_payload() or '')
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_r_1_1_happy_path(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
):
|
||||
"""R.1.1: Mixed-level notifications produce one email with all sections."""
|
||||
seed_receiver_groups(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[default_receiver_group(contents=['reports'], levels=['ERROR', 'WARNING', 'INFO'])],
|
||||
)
|
||||
seed_notifications(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[
|
||||
default_notification(notification_id='r-err', level='ERROR', model_name='m1'),
|
||||
default_notification(
|
||||
notification_id='r-warn',
|
||||
level='WARNING',
|
||||
model_name='m2',
|
||||
timestamp='2024-06-01 12:01:00+0000',
|
||||
),
|
||||
default_notification(
|
||||
notification_id='r-info',
|
||||
level='INFO',
|
||||
model_name='m3',
|
||||
timestamp='2024-06-01 12:02:00+0000',
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Reports.run,
|
||||
load_scenario_input('reports_happy_path.json'),
|
||||
make_workflow_id('reports-happy'),
|
||||
)
|
||||
|
||||
assert len(smtp_server.messages) == 1
|
||||
body = _extract_html_body(smtp_server.messages[0])
|
||||
assert 'Errors detected:' in body
|
||||
assert 'Warnings detected:' in body
|
||||
assert 'Infos detected:' in body
|
||||
# format_log_report writes one row per (notification_id, trigger) pair, so the
|
||||
# three seeded notifications produce three rows even though a single email
|
||||
# was sent to the receiver group.
|
||||
assert count_log_report_rows(postgres_engine, 'Reports') == 3
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_r_1_2_error_section_only(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
):
|
||||
"""R.1.2: Only ERROR notifications → only Errors section in HTML."""
|
||||
seed_receiver_groups(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[default_receiver_group(contents=['reports'], levels=['ERROR'])],
|
||||
)
|
||||
seed_notifications(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[default_notification(notification_id='only-err', level='ERROR')],
|
||||
)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Reports.run,
|
||||
load_scenario_input('reports_multi_level.json'),
|
||||
make_workflow_id('reports-error-only'),
|
||||
)
|
||||
|
||||
assert smtp_server.messages, 'Expected at least one report email'
|
||||
body = _extract_html_body(smtp_server.messages[0])
|
||||
assert 'Errors detected:' in body
|
||||
assert 'Warnings detected:' not in body
|
||||
assert 'Infos detected:' not in body
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_scenario_r_2_1_empty_queue(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
):
|
||||
"""R.2.1: Empty queue → no email and no log_report row."""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
Reports.run,
|
||||
load_scenario_input('reports_empty.json'),
|
||||
make_workflow_id('reports-empty'),
|
||||
)
|
||||
|
||||
assert len(smtp_server.messages) == 0
|
||||
assert count_log_report_rows(postgres_engine, 'Reports') == 0
|
||||
154
e2e/test_subworkflow_load_notification_package.py
Normal file
154
e2e/test_subworkflow_load_notification_package.py
Normal file
@@ -0,0 +1,154 @@
|
||||
"""E2E tests for LoadNotificationPackage subworkflow."""
|
||||
|
||||
import pytest
|
||||
from redis import Redis
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
|
||||
from e2e.conftest import E2E_DATABASE
|
||||
from e2e.helpers import (
|
||||
default_notification,
|
||||
default_receiver_group,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
seed_last_timestamp,
|
||||
seed_notifications,
|
||||
seed_receiver_groups,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
|
||||
|
||||
|
||||
def _build_metadata(input_data: dict) -> None:
|
||||
"""Attach the metadata block that the subworkflow expects."""
|
||||
input_data['metadata'] = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'workflow_name': 'load_notification_package',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_package_without_prior_timestamp(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""
|
||||
No Redis timestamp → returns notifications and stores max timestamp.
|
||||
|
||||
Notifications are seeded with the exact production format produced by
|
||||
``sientia_do.notifications.models.Notification`` (string in
|
||||
``DATETIME_FORMAT_WITH_TZ``, e.g. ``"2024-06-01 11:00:00+0000"``). The
|
||||
cached "last timestamp" must mirror that representation.
|
||||
"""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
seed_notifications(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[
|
||||
default_notification(
|
||||
notification_id='n1',
|
||||
timestamp='2024-06-01 10:00:00+0000',
|
||||
),
|
||||
default_notification(
|
||||
notification_id='n2',
|
||||
timestamp='2024-06-01 11:00:00+0000',
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
input_data = load_scenario_input('subworkflow_load_notification_package.json')
|
||||
_build_metadata(input_data)
|
||||
|
||||
result = await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
LoadNotificationPackage.run,
|
||||
input_data,
|
||||
make_workflow_id('load-pkg-none'),
|
||||
)
|
||||
|
||||
assert len(result['notification_package']) == 2
|
||||
stored = redis_client.get('notification_last_timestamp:Alerts') or ''
|
||||
assert stored.strip('"') == '2024-06-01 11:00:00+0000'
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_package_with_prior_timestamp(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""
|
||||
Existing timestamp → only newer notifications returned.
|
||||
|
||||
Both the seeded "last timestamp" (Redis) and the notification timestamps
|
||||
(Mongo) follow the production format used by
|
||||
``sientia_do.notifications.models.Notification`` (string in
|
||||
``DATETIME_FORMAT_WITH_TZ``). ``load_latest_data`` will translate the
|
||||
Redis value into a Python ``datetime`` and apply ``{$gt: <Date>}``
|
||||
against the Mongo string timestamps; this exercise surfaces the real
|
||||
BSON comparison semantics rather than a synthetic ideal.
|
||||
"""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
seed_last_timestamp(redis_client, 'Alerts', '2024-06-01 10:30:00+0000')
|
||||
seed_notifications(
|
||||
mongo_uri,
|
||||
E2E_DATABASE,
|
||||
[
|
||||
default_notification(
|
||||
notification_id='old',
|
||||
timestamp='2024-06-01 10:00:00+0000',
|
||||
),
|
||||
default_notification(
|
||||
notification_id='new',
|
||||
timestamp='2024-06-01 11:00:00+0000',
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
input_data = load_scenario_input('subworkflow_load_notification_package.json')
|
||||
_build_metadata(input_data)
|
||||
|
||||
result = await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
LoadNotificationPackage.run,
|
||||
input_data,
|
||||
make_workflow_id('load-pkg-ts'),
|
||||
)
|
||||
|
||||
ids = {n['notification_id'] for n in result['notification_package']}
|
||||
assert ids == {'new'}
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_package_empty_mongo_no_redis_write(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
mongo_uri: str,
|
||||
redis_client: Redis,
|
||||
):
|
||||
"""Empty Mongo → no Redis timestamp write."""
|
||||
seed_receiver_groups(mongo_uri, E2E_DATABASE, [default_receiver_group()])
|
||||
|
||||
input_data = load_scenario_input('subworkflow_load_notification_package.json')
|
||||
_build_metadata(input_data)
|
||||
|
||||
await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
LoadNotificationPackage.run,
|
||||
input_data,
|
||||
make_workflow_id('load-pkg-empty'),
|
||||
)
|
||||
|
||||
assert redis_client.get('notification_last_timestamp:Alerts') is None
|
||||
94
e2e/test_subworkflow_process_notifications.py
Normal file
94
e2e/test_subworkflow_process_notifications.py
Normal file
@@ -0,0 +1,94 @@
|
||||
"""E2E tests for ProcessNotifications subworkflow."""
|
||||
|
||||
import pytest
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
|
||||
from e2e.helpers import (
|
||||
count_log_report_rows,
|
||||
default_notification,
|
||||
fetch_log_report,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
|
||||
|
||||
|
||||
def _receiver_package():
|
||||
notif = default_notification(notification_id='proc-1', level='ERROR')
|
||||
return {
|
||||
'admins': {
|
||||
'group_name': 'admins',
|
||||
'members': ['admin@example.com'],
|
||||
'status': 'pending',
|
||||
'notifications': [notif],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_notifications_round_trip(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
):
|
||||
"""Email → SMTP → log_report persisted in Postgres."""
|
||||
input_data = load_scenario_input('subworkflow_process_notifications.json')
|
||||
input_data['metadata'] = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'workflow_name': 'process_notifications',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
input_data['notification_package'] = _receiver_package()
|
||||
|
||||
result = await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
ProcessNotifications.run,
|
||||
input_data,
|
||||
make_workflow_id('process-notif'),
|
||||
)
|
||||
|
||||
assert result
|
||||
assert len(smtp_server.messages) == 1
|
||||
assert count_log_report_rows(postgres_engine, 'Alerts') == 1
|
||||
rows = fetch_log_report(postgres_engine, 'Alerts')
|
||||
assert rows[0]['notification_id'] == 'proc-1'
|
||||
|
||||
|
||||
@pytest.mark.e2e
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_notifications_empty_package(
|
||||
temporal_env: WorkflowEnvironment,
|
||||
orchestrator_worker,
|
||||
stub_workers,
|
||||
smtp_server,
|
||||
postgres_engine,
|
||||
):
|
||||
"""Empty receiver groups returns empty dict."""
|
||||
input_data = load_scenario_input('subworkflow_process_notifications.json')
|
||||
input_data['metadata'] = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'workflow_name': 'process_notifications',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
input_data['notification_package'] = {}
|
||||
|
||||
result = await start_and_await_workflow(
|
||||
temporal_env.client,
|
||||
ProcessNotifications.run,
|
||||
input_data,
|
||||
make_workflow_id('process-empty'),
|
||||
)
|
||||
|
||||
assert result == {}
|
||||
assert len(smtp_server.messages) == 0
|
||||
assert count_log_report_rows(postgres_engine) == 0
|
||||
113
encrypt.py
Normal file
113
encrypt.py
Normal file
@@ -0,0 +1,113 @@
|
||||
import os
|
||||
import argparse
|
||||
from pathspec import PathSpec
|
||||
import yaml # type: ignore
|
||||
from typing import Any
|
||||
|
||||
'''
|
||||
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: dict[str, Any] = {}
|
||||
|
||||
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()
|
||||
1
git-requirements-mapping.txt
Normal file
1
git-requirements-mapping.txt
Normal file
@@ -0,0 +1 @@
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia-do
|
||||
234
init_orchestration.ipynb
Normal file
234
init_orchestration.ipynb
Normal file
File diff suppressed because one or more lines are too long
416
notification_generator.ipynb
Normal file
416
notification_generator.ipynb
Normal file
File diff suppressed because one or more lines are too long
1
orchestrator/__init__.py
Normal file
1
orchestrator/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
|
||||
0
orchestrator/activities/__init__.py
Normal file
0
orchestrator/activities/__init__.py
Normal file
130
orchestrator/activities/activities.py
Normal file
130
orchestrator/activities/activities.py
Normal file
@@ -0,0 +1,130 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from logging import Logger
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.temporal.activities.postgres_sync import Postgres
|
||||
|
||||
from orchestrator.activities.email import Email
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
from orchestrator.activities.mongo_db import MongoDB
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
|
||||
|
||||
class Activities(TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres):
|
||||
"""
|
||||
Central activities orchestrator for Temporal workflow operations.
|
||||
|
||||
This class combines multiple activity components including temporal management,
|
||||
slot management, data formatting, MongoDB operations, email services, and
|
||||
PostgreSQL operations. It provides a unified interface for all activity
|
||||
operations required by the orchestration workflows.
|
||||
|
||||
Args:
|
||||
temporal_config (dict[str, Any]): Temporal server configuration
|
||||
redis_config (dict[str, Any]): Redis server configuration
|
||||
mongodb_config (dict[str, Any]): MongoDB connection configuration
|
||||
email_config (dict[str, Any]): Email service configuration
|
||||
postgres_config (dict[str, Any]): PostgreSQL database configuration
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
temporal_config: dict[str, Any],
|
||||
# couchbase_config: dict[str, Any],
|
||||
redis_config: dict[str, Any],
|
||||
mongodb_config: dict[str, Any],
|
||||
email_config: dict[str, Any],
|
||||
postgres_config: dict[str, Any],
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
):
|
||||
# Initialize parent classes
|
||||
|
||||
metrics_controller = MetricsController(logger=logger)
|
||||
|
||||
TemporalManager.__init__(
|
||||
self,
|
||||
host=temporal_config['temporal_host'],
|
||||
scouter_namespace=temporal_config['temporal_scouter_namespace'],
|
||||
laborious_namespace=temporal_config['temporal_laborious_namespace'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
SlotManager.__init__(
|
||||
self,
|
||||
host=redis_config['host'],
|
||||
port=redis_config['port'],
|
||||
username=redis_config['username'],
|
||||
password=redis_config['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
Formatters.__init__(
|
||||
self,
|
||||
scouter_namespace=temporal_config['temporal_scouter_namespace'],
|
||||
laborious_namespace=temporal_config['temporal_laborious_namespace'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
MongoDB.__init__(
|
||||
self,
|
||||
connection_string=mongodb_config['connection_string'],
|
||||
database_name=mongodb_config['database_name'],
|
||||
ttl_index_seconds=mongodb_config['ttl_index_seconds'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
Email.__init__(
|
||||
self,
|
||||
sender_email=email_config['sender_email'],
|
||||
sender_password=email_config['sender_password'],
|
||||
smtp_server=email_config['smtp_server'],
|
||||
smtp_port=email_config['smtp_port'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
Postgres.__init__(
|
||||
self,
|
||||
host=postgres_config['host'],
|
||||
port=postgres_config['port'],
|
||||
user=postgres_config['user'],
|
||||
password=postgres_config['password'],
|
||||
dbname=postgres_config['dbname'],
|
||||
min_connections=postgres_config['min_connections'],
|
||||
max_connections=postgres_config['max_connections'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def shutdown(self):
|
||||
"""
|
||||
Shutdown all connections and clean up resources.
|
||||
|
||||
This method gracefully shuts down all database connections, email
|
||||
services, and other resources to ensure proper cleanup when the
|
||||
application terminates.
|
||||
"""
|
||||
MongoDB.close(self)
|
||||
Postgres.close(self)
|
||||
Email.close(self)
|
||||
Formatters.close(self)
|
||||
SlotManager.close(self)
|
||||
TemporalManager.close(self)
|
||||
269
orchestrator/activities/email.py
Normal file
269
orchestrator/activities/email.py
Normal file
@@ -0,0 +1,269 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import smtplib
|
||||
import traceback
|
||||
from email import encoders
|
||||
from email.mime.base import MIMEBase
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.text import MIMEText
|
||||
from smtplib import SMTPServerDisconnected
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
|
||||
|
||||
from orchestrator import metrics
|
||||
from orchestrator.utils.email_builder import EmailBuilder
|
||||
|
||||
|
||||
class Email(SientiaMonitoring):
|
||||
"""
|
||||
Email service activity for sending workflow notifications.
|
||||
|
||||
This class provides email sending capabilities including HTML email
|
||||
generation, attachment handling, and SMTP connection management with
|
||||
automatic reconnection for workflow notification delivery.
|
||||
|
||||
Args:
|
||||
sender_email (str): Email address for sending messages
|
||||
sender_password (str): SMTP authentication password
|
||||
smtp_server (str): SMTP server hostname
|
||||
smtp_port (int): SMTP server port number
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sender_email: str,
|
||||
sender_password: str,
|
||||
smtp_server: str,
|
||||
smtp_port: int,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
self.email_builder = EmailBuilder(logger=logger)
|
||||
|
||||
self.sender_email = sender_email
|
||||
self.sender_password = sender_password
|
||||
self.smtp_port = smtp_port
|
||||
self.smtp_server = smtp_server
|
||||
|
||||
logger.info(f'Initializing Email with {smtp_server}:{smtp_port}')
|
||||
|
||||
if smtp_server is not None:
|
||||
self.server = smtplib.SMTP(smtp_server, smtp_port, timeout=20)
|
||||
|
||||
if self.sender_password:
|
||||
self.server.starttls()
|
||||
self.server.login(self.sender_email, self.sender_password)
|
||||
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Close the Email connection and clean up resources.
|
||||
|
||||
Closes the SMTP server connection and shuts down the SientiaMonitoring instance.
|
||||
"""
|
||||
self.server.quit()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Ensure the Email connection is closed when the object is garbage-collected.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='build_email_html')
|
||||
def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Build HTML email content for configured receiver groups.
|
||||
|
||||
This activity generates HTML email content for each receiver group
|
||||
based on notification data and mail type. It processes notification
|
||||
data through the email builder to create formatted HTML messages.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Activity input parameters.
|
||||
Required fields:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- receiver_groups (dict[str, Any]): Receiver group configurations with notifications
|
||||
- mail_type (str): Type of email (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Updated receiver groups with generated HTML content
|
||||
|
||||
Raises:
|
||||
Exception: If HTML generation fails
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
receiver_groups = input_data['receiver_groups']
|
||||
mail_type = input_data['mail_type']
|
||||
|
||||
self.info(f'Building email html for {mail_type} mail type.', metadata=metadata)
|
||||
|
||||
for _group_name, group_config in receiver_groups.items():
|
||||
html = self.email_builder.build_email(group_config['notifications'], mail_type)
|
||||
|
||||
group_config['html'] = html
|
||||
|
||||
self.info(f'Email html built for {mail_type} mail type.', metadata=metadata)
|
||||
|
||||
return receiver_groups
|
||||
|
||||
def handle_attachments(self, attachments: list[dict], msg: MIMEMultipart) -> MIMEMultipart:
|
||||
"""
|
||||
Attach a list of attachments to an email message.
|
||||
|
||||
Processes notification attachments and adds them to the email message
|
||||
as base64-encoded MIME parts. Each attachment contains error details
|
||||
or additional context for the notification.
|
||||
|
||||
Args:
|
||||
attachments (list[dict]): A list of dictionaries where each dictionary contains:
|
||||
- filename (str): Name of the attachment file
|
||||
- attachment_content (str): Content of the attachment
|
||||
msg (MIMEMultipart): The email message object to which the attachments will be added
|
||||
|
||||
Returns:
|
||||
MIMEMultipart: The email message object with the attachments added
|
||||
|
||||
Raises:
|
||||
Exception: If an attachment cannot be added, an error is logged and raised
|
||||
"""
|
||||
|
||||
for attachment in attachments:
|
||||
att_name = attachment['filename']
|
||||
try:
|
||||
# Create the attachment as a MIMEBase object
|
||||
part = MIMEBase('application', 'octet-stream')
|
||||
part.set_payload(attachment['attachment_content'].encode('utf-8'))
|
||||
encoders.encode_base64(part)
|
||||
part.add_header('Content-Disposition', f'attachment; filename="{att_name}"')
|
||||
msg.attach(part)
|
||||
except Exception as e:
|
||||
self.logger.error(f'Failed to attach content of {att_name}: {e}')
|
||||
|
||||
raise e
|
||||
|
||||
return msg
|
||||
|
||||
def try_send_email(self, msg: MIMEMultipart, receivers: str):
|
||||
"""
|
||||
Send an email to the receivers with automatic reconnection handling.
|
||||
|
||||
Attempts to send an email and automatically reconnects to the SMTP server
|
||||
if a disconnection occurs during transmission.
|
||||
|
||||
Args:
|
||||
msg (MIMEMultipart): The email message to send
|
||||
receivers (str): Comma-separated list of email addresses to send to
|
||||
|
||||
Raises:
|
||||
Exception: If email sending fails after reconnection attempts
|
||||
"""
|
||||
try:
|
||||
self.server.sendmail(self.sender_email, receivers, msg.as_string())
|
||||
except SMTPServerDisconnected as e:
|
||||
self.logger.error(f'SMTP server disconnected: {e}')
|
||||
self.logger.info(f'Reconnecting to {self.smtp_server}:{self.smtp_port}')
|
||||
|
||||
if self.server:
|
||||
try:
|
||||
self.server.quit()
|
||||
except SMTPServerDisconnected as e:
|
||||
self.logger.info(f'Server already disconnected: {e}')
|
||||
except Exception as e:
|
||||
self.logger.error(f'Failed to quit server: {e}')
|
||||
raise e
|
||||
|
||||
self.server = smtplib.SMTP(self.smtp_server, self.smtp_port, timeout=20)
|
||||
if self.sender_password:
|
||||
self.server.starttls()
|
||||
self.server.login(self.sender_email, self.sender_password)
|
||||
self.server.sendmail(self.sender_email, receivers, msg.as_string())
|
||||
|
||||
@activity.defn(name='send_email')
|
||||
def send_email(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Send email notifications to configured receiver groups.
|
||||
|
||||
This activity sends HTML emails with attachments to all configured
|
||||
receiver groups. It handles SMTP connection management, attachment
|
||||
processing, and error reporting with automatic reconnection support.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Activity input parameters.
|
||||
Required fields:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- receiver_groups (dict[str, Any]): Receiver groups with HTML content
|
||||
- mail_type (str): Type of email being sent (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Updated receiver groups with sending status
|
||||
|
||||
Raises:
|
||||
Exception: If email sending fails for all groups
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
receiver_groups = input_data['receiver_groups']
|
||||
mail_type = input_data['mail_type']
|
||||
|
||||
if self.smtp_server is None:
|
||||
self.info(f'Skipping email sending for {mail_type} mail type.', metadata=metadata)
|
||||
return {}
|
||||
|
||||
self.info(f'Sending email for {mail_type} mail type.', metadata=metadata)
|
||||
|
||||
for group_name, group_config in receiver_groups.items():
|
||||
try:
|
||||
receivers = ', '.join(group_config['members'])
|
||||
|
||||
self.info(f'Sending email to {group_name}: {receivers}', metadata=metadata)
|
||||
|
||||
msg = MIMEMultipart()
|
||||
msg.attach(MIMEText(group_config['html'], 'html'))
|
||||
msg['From'] = self.sender_email
|
||||
msg['To'] = receivers
|
||||
msg['Subject'] = f'SIENTIA™ {mail_type}'
|
||||
|
||||
msg = self.handle_attachments(
|
||||
[
|
||||
{
|
||||
'filename': f'{notification["trigger"]}_{notification["notification_id"]}.txt',
|
||||
'attachment_content': notification['attachment_content'],
|
||||
}
|
||||
for notification in group_config['notifications']
|
||||
if notification.get('attachment_content') is not None
|
||||
],
|
||||
msg,
|
||||
)
|
||||
|
||||
self.try_send_email(msg, receivers)
|
||||
except Exception as e:
|
||||
self.error(f'Failed to send email to {group_name}: {e}', metadata=metadata)
|
||||
traceback.print_exc()
|
||||
group_config['status'] = 'failed'
|
||||
else:
|
||||
group_config['status'] = 'sent'
|
||||
metrics.EMAIL_SENT_COUNT.labels(
|
||||
pod_id=self.pod_id,
|
||||
model_name=metadata['model_name'],
|
||||
pipeline_name=metadata['workflow_name'],
|
||||
email_group=group_name,
|
||||
).inc()
|
||||
|
||||
self.info(f'Email sent to {group_name}: {receivers}', metadata=metadata)
|
||||
|
||||
self.info(f'Email sent for {mail_type} mail type.', metadata=metadata)
|
||||
|
||||
return receiver_groups
|
||||
854
orchestrator/activities/formatters.py
Normal file
854
orchestrator/activities/formatters.py
Normal file
@@ -0,0 +1,854 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from logging import Logger
|
||||
from math import ceil
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
from pandas import DataFrame
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
|
||||
|
||||
from orchestrator.utils.orchestrator_functions import (
|
||||
build_tag_config,
|
||||
drift,
|
||||
gather_read_tags,
|
||||
minimal_retrain,
|
||||
pi_web_api_scouter,
|
||||
predictions_batch,
|
||||
scouter,
|
||||
simple_metrics,
|
||||
)
|
||||
|
||||
topic_separator = '\n ========== \n'
|
||||
|
||||
|
||||
class ScheduleType(TypedDict):
|
||||
"""
|
||||
Type definition for schedule type configuration entries.
|
||||
|
||||
Each schedule type entry maps a workflow type string to its corresponding
|
||||
namespace and configuration builder function.
|
||||
|
||||
Attributes:
|
||||
namespace (str): The Temporal namespace where workflows of this type execute.
|
||||
Valid values: 'scouter', 'laborious'
|
||||
function (Callable): Configuration builder function that transforms pipeline
|
||||
configuration into Temporal-compatible workflow arguments
|
||||
"""
|
||||
|
||||
namespace: str
|
||||
function: Callable
|
||||
|
||||
|
||||
schedule_types: dict[str, ScheduleType] = {
|
||||
'scouter': {
|
||||
'namespace': 'scouter',
|
||||
'function': scouter,
|
||||
},
|
||||
'pi_web_api_scouter': {
|
||||
'namespace': 'scouter',
|
||||
'function': pi_web_api_scouter,
|
||||
},
|
||||
'predictions_batch': {
|
||||
'namespace': 'laborious',
|
||||
'function': predictions_batch,
|
||||
},
|
||||
'xgboost_predictions_batch': {
|
||||
'namespace': 'laborious',
|
||||
'function': predictions_batch,
|
||||
},
|
||||
'minimal_retrain': {
|
||||
'namespace': 'laborious',
|
||||
'function': minimal_retrain,
|
||||
},
|
||||
'xgboost_minimal_retrain': {
|
||||
'namespace': 'laborious',
|
||||
'function': minimal_retrain,
|
||||
},
|
||||
'drift': {
|
||||
'namespace': 'laborious',
|
||||
'function': drift,
|
||||
},
|
||||
'simple_metrics': {
|
||||
'namespace': 'laborious',
|
||||
'function': simple_metrics,
|
||||
},
|
||||
}
|
||||
"""
|
||||
Registry mapping workflow types to their namespace and configuration builder functions.
|
||||
|
||||
Supported workflow types:
|
||||
- scouter: OPC data collection using OPC UA protocol
|
||||
- pi_web_api_scouter: Data collection using PI Web API
|
||||
- predictions_batch: ML model prediction workflows with OPC write-back
|
||||
- minimal_retrain: Model retraining workflows using SQL queries
|
||||
- drift: Data drift detection and monitoring workflows
|
||||
- simple_metrics: Model performance metrics computation workflows
|
||||
|
||||
Each entry specifies:
|
||||
- namespace: Target Temporal namespace for workflow execution
|
||||
- function: Configuration builder that transforms MongoDB pipeline config
|
||||
into Temporal workflow arguments
|
||||
"""
|
||||
|
||||
|
||||
class Formatters(SientiaMonitoring):
|
||||
"""
|
||||
Schedule and slot configuration formatting and notification filtering activity.
|
||||
|
||||
This class provides comprehensive formatting operations for schedules and OPC slots,
|
||||
converting pipeline configurations into Temporal-compatible formats, managing slot
|
||||
distribution across active ingestors, and implementing notification filtering for
|
||||
scheduled reports.
|
||||
|
||||
Key features:
|
||||
- Pipeline schedule configuration formatting ("scouter", "predictions_batch", "minimal_retrain", "drift")
|
||||
- OPC slot distribution across active ingestors
|
||||
- Notification filtering for comprehensive scheduled reports
|
||||
- Group-based report filtering with ignore list support
|
||||
- Configuration validation and transformation
|
||||
|
||||
Args:
|
||||
scouter_namespace (str): Scouter workflow namespace
|
||||
laborious_namespace (str): Laborious workflow namespace
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
scouter_namespace: str,
|
||||
laborious_namespace: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
self.scouter_namespace = scouter_namespace
|
||||
self.laborious_namespace = laborious_namespace
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Close the Formatters connection and clean up resources.
|
||||
|
||||
Shuts down the SientiaMonitoring instance and releases all resources.
|
||||
"""
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Ensure the Formatters connection is closed when the object is garbage-collected.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='process_schedules')
|
||||
def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Process pipeline configurations into Temporal-compatible schedule configurations.
|
||||
|
||||
This method transforms pipeline configurations from MongoDB into properly formatted
|
||||
Temporal schedule configurations, organizing them by workflow type (scouter and
|
||||
laborious) and applying the appropriate configuration builders for each pipeline type.
|
||||
|
||||
Pipeline types supported:
|
||||
- "scouter": Data collection workflows with OPC tag configurations
|
||||
- "predictions_batch": ML prediction workflows with OPC write configurations
|
||||
- "minimal_retrain": Model retraining workflows with SQL query configurations
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to process.
|
||||
- pipelines (list[dict[str, Any]]): The schedules to process.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: The schedule configuration dictionary keyed by namespace
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Processing schedules...', metadata=metadata)
|
||||
|
||||
pipelines = input_data['pipelines']
|
||||
|
||||
schedule_config: dict[str, dict[str, Any]] = {
|
||||
self.scouter_namespace: {},
|
||||
self.laborious_namespace: {},
|
||||
}
|
||||
|
||||
for pipeline in pipelines:
|
||||
workflow_type = pipeline['workflow_type']
|
||||
|
||||
if workflow_type not in schedule_types:
|
||||
self.error(f'Workflow type {workflow_type} not supported', metadata=metadata)
|
||||
continue
|
||||
|
||||
schedule_type = schedule_types[workflow_type]
|
||||
namespace = schedule_type['namespace']
|
||||
function = schedule_type['function']
|
||||
|
||||
schedule_config[namespace][pipeline['schedule_name']] = {
|
||||
**function(pipeline),
|
||||
'updated_at': pipeline.get(
|
||||
'updated_at', now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
|
||||
),
|
||||
}
|
||||
|
||||
self.info('Processed schedules', metadata=metadata)
|
||||
self.debug(json.dumps(schedule_config, indent=4, sort_keys=True), metadata=metadata)
|
||||
|
||||
return schedule_config
|
||||
|
||||
@activity.defn(name='process_slots')
|
||||
def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Extracts all read tags from input pipelines, divides them into slots and
|
||||
returns a slot config dictionary. If no ingestor is available, only one slot
|
||||
is created.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to process.
|
||||
- pipelines (list[dict[str, Any]]): The schedules to process.
|
||||
- opc_servers (list[str]): The OPC servers to create ingestor config.
|
||||
- active_ingestors (list[str]): The active ingestors to divide into slots.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: The slot configuration dictionary keyed by slot id (as string)
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Processing slots...', metadata=metadata)
|
||||
|
||||
pipelines = input_data['pipelines']
|
||||
opc_servers_list = input_data['opc_servers']
|
||||
active_ingestors = input_data['active_ingestors']
|
||||
|
||||
opc_servers = {}
|
||||
for server in opc_servers_list:
|
||||
opc_servers[server['id']] = {
|
||||
**server,
|
||||
}
|
||||
|
||||
tags = list(gather_read_tags(pipelines).values())
|
||||
|
||||
number_of_tags = len(tags)
|
||||
number_of_slots = len(active_ingestors) if active_ingestors else 1
|
||||
tags_per_slot = ceil(number_of_tags / number_of_slots)
|
||||
|
||||
slot_config: dict[str, Any] = {}
|
||||
last_index = 0
|
||||
|
||||
for i in range(1, number_of_slots):
|
||||
slot_tags = tags[last_index : last_index + tags_per_slot]
|
||||
slot_config[f'{i}'], notifications = build_tag_config(slot_tags, opc_servers)
|
||||
|
||||
if notifications:
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
|
||||
message=f'Servers {", ".join(notifications)} not found in opc_servers',
|
||||
block='orchestrator',
|
||||
level=NotificationLevel.ERROR,
|
||||
)
|
||||
|
||||
last_index += tags_per_slot
|
||||
|
||||
slot_tags = tags[last_index:]
|
||||
slot_config[f'{number_of_slots}'], notifications = build_tag_config(slot_tags, opc_servers)
|
||||
if notifications:
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
|
||||
message=f'Servers {", ".join(notifications)} not found in opc_servers',
|
||||
block='orchestrator',
|
||||
level=NotificationLevel.ERROR,
|
||||
)
|
||||
|
||||
self.info('Processed slots', metadata=metadata)
|
||||
self.debug(json.dumps(slot_config, indent=4, sort_keys=True), metadata=metadata)
|
||||
|
||||
return slot_config
|
||||
|
||||
@activity.defn(name='format_schedule_config')
|
||||
def format_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Format the schedule config to a dictionary with the schedule name as the key.
|
||||
|
||||
Transforms a list of schedule configurations into a nested dictionary structure
|
||||
organized by namespace and schedule name for efficient lookup and comparison.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- schedule_config (list[dict[str, Any]]): The schedule config to format
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: The formatted schedule config organized by namespace and schedule name
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
self.info('Formatting schedule config...', metadata=metadata)
|
||||
|
||||
schedule_config = input_data['schedule_config']
|
||||
|
||||
config: dict[str, dict[str, str]] = {}
|
||||
for schedule in schedule_config:
|
||||
namespace = schedule['namespace']
|
||||
schedule_name = schedule['schedule_name']
|
||||
updated_at = schedule['updated_at']
|
||||
|
||||
if namespace not in config:
|
||||
config[namespace] = {}
|
||||
|
||||
config[namespace][schedule_name] = updated_at
|
||||
|
||||
self.info('Formatted schedule config', metadata=metadata)
|
||||
self.debug(json.dumps(config, indent=4, sort_keys=True), metadata=metadata)
|
||||
|
||||
return config
|
||||
|
||||
def compare_config_timestamps(
|
||||
self,
|
||||
schedules: dict[str, Any],
|
||||
current_schedules: dict[str, Any],
|
||||
to_update: dict[str, Any],
|
||||
to_create: dict[str, Any],
|
||||
namespace: str,
|
||||
metadata: dict[str, Any],
|
||||
):
|
||||
"""
|
||||
Compare new and current schedules to determine which should be updated or created.
|
||||
|
||||
Compares schedule timestamps to identify schedules that need updating (newer timestamp)
|
||||
or creating (schedule doesn't exist). Results are accumulated in the provided dictionaries.
|
||||
|
||||
Args:
|
||||
schedules (dict[str, Any]): New schedules to compare
|
||||
current_schedules (dict[str, Any]): Existing schedules to compare against
|
||||
to_update (dict[str, Any]): Output accumulator for schedules that need updating
|
||||
to_create (dict[str, Any]): Output accumulator for schedules that need creating
|
||||
namespace (str): Namespace for the schedules being compared
|
||||
metadata (dict[str, Any]): Metadata for logging purposes
|
||||
"""
|
||||
for schedule_name, schedule in schedules.items():
|
||||
if schedule_name in current_schedules:
|
||||
update_timestamp = schedule.get('updated_at', now())
|
||||
|
||||
old_timestamp = current_schedules[schedule_name]
|
||||
|
||||
self.debug(
|
||||
f'Comparing schedule {schedule_name}:{update_timestamp} vs {old_timestamp}',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
if update_timestamp > old_timestamp:
|
||||
to_update[namespace][schedule_name] = schedule
|
||||
else:
|
||||
to_create[namespace][schedule_name] = schedule
|
||||
|
||||
@activity.defn(name='create_schedule_config')
|
||||
def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Create a schedule config dictionary based on the input data.
|
||||
|
||||
Compares the current schedule configuration in Temporal with the new schedule
|
||||
configuration to determine which schedules need to be created, updated, or deleted.
|
||||
Uses timestamp comparison to identify schedules that have changed.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- current_schedule_config (dict[str, Any]): The current schedule config in Temporal server
|
||||
- schedule_config (dict[str, Any]): The schedule config to process
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: A dictionary with keys:
|
||||
- to_update (dict): Schedules that need updating
|
||||
- to_create (dict): Schedules that need creating
|
||||
- to_delete (dict): Schedules that need deleting
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Creating schedule config...', metadata=metadata)
|
||||
|
||||
current_schedule_config = input_data['current_schedule_config']
|
||||
schedule_config = input_data['schedule_config']
|
||||
|
||||
to_update: dict[str, dict[str, Any]] = {
|
||||
self.scouter_namespace: {},
|
||||
self.laborious_namespace: {},
|
||||
}
|
||||
to_create: dict[str, dict[str, Any]] = {
|
||||
self.scouter_namespace: {},
|
||||
self.laborious_namespace: {},
|
||||
}
|
||||
to_delete: dict[str, list[str]] = {
|
||||
self.scouter_namespace: [],
|
||||
self.laborious_namespace: [],
|
||||
}
|
||||
|
||||
for namespace, schedules in schedule_config.items():
|
||||
current_schedules = current_schedule_config.get(namespace, {})
|
||||
self.compare_config_timestamps(
|
||||
schedules, current_schedules, to_update, to_create, namespace, metadata
|
||||
)
|
||||
|
||||
for namespace, schedules in current_schedule_config.items():
|
||||
for schedule_name in schedules:
|
||||
if schedule_name not in schedule_config[namespace]:
|
||||
to_delete[namespace].append(schedule_name)
|
||||
|
||||
output = {'to_update': to_update, 'to_create': to_create, 'to_delete': to_delete}
|
||||
|
||||
self.info('Created schedule config', metadata=metadata)
|
||||
self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
|
||||
|
||||
return output
|
||||
|
||||
@activity.defn(name='create_slot_config')
|
||||
def create_slot_config(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Create a slot config dictionary based on the input data.
|
||||
|
||||
Compares the current slot configuration in Redis with the new slot configuration
|
||||
to determine which slots need to be inserted or deleted. Slots are identified
|
||||
by numeric IDs, and excess slots are marked for deletion.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- current_slot_config (dict[str, Any]): The current slot config in Redis
|
||||
- slot_config (dict[str, Any]): The slot config to process
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: A dictionary containing:
|
||||
- to_insert (dict): Slots that need to be inserted/updated
|
||||
- to_delete (list[str]): Slot IDs that need to be deleted
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Creating slot config...', metadata=metadata)
|
||||
|
||||
current_slot_config = input_data['current_slot_config']
|
||||
slot_config = input_data['slot_config']
|
||||
to_delete = []
|
||||
|
||||
number_of_current_slots = len(current_slot_config)
|
||||
number_of_slots = len(slot_config)
|
||||
|
||||
if number_of_current_slots > number_of_slots:
|
||||
to_delete = [str(i) for i in range(number_of_slots + 1, number_of_current_slots + 1)]
|
||||
|
||||
output = {'to_delete': to_delete, 'to_insert': slot_config}
|
||||
|
||||
self.info('Created slot config', metadata=metadata)
|
||||
self.debug(json.dumps(output, indent=4, sort_keys=True), metadata=metadata)
|
||||
|
||||
return output
|
||||
|
||||
def send_success_report(
|
||||
self,
|
||||
metadata: dict[str, Any],
|
||||
message: str,
|
||||
notification_id: str,
|
||||
attachment: Any | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Send a success notification report.
|
||||
|
||||
Sends an INFO-level notification to the notification handler with success
|
||||
details about orchestration operations.
|
||||
|
||||
Args:
|
||||
metadata (dict[str, Any]): Metadata for the notification
|
||||
message (str): The success message to send
|
||||
notification_id (str): The ID of the notification to send
|
||||
attachment (Any | None, optional): Optional attachment content to include
|
||||
"""
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=notification_id,
|
||||
message=message,
|
||||
block='report_orchestration',
|
||||
level=NotificationLevel.INFO,
|
||||
attachment_content=json.dumps(attachment, indent=4, sort_keys=True),
|
||||
)
|
||||
|
||||
def send_error_report(
|
||||
self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str
|
||||
) -> None:
|
||||
"""
|
||||
Send an error notification report.
|
||||
|
||||
Sends an ERROR-level notification to the notification handler with error
|
||||
details about orchestration operation failures.
|
||||
|
||||
Args:
|
||||
metadata (dict[str, Any]): Metadata for the notification
|
||||
message (str): The error message to send
|
||||
notification_id (str): The ID of the notification
|
||||
attachment (str): The attachment content for the notification
|
||||
"""
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=notification_id,
|
||||
message=message,
|
||||
block='report_orchestration',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=attachment,
|
||||
)
|
||||
|
||||
def parse_report_schedule(
|
||||
self, input_data: list[dict[str, Any]]
|
||||
) -> tuple[list[str], dict[str, Any]]:
|
||||
"""
|
||||
Parse the report schedule data to extract success and error information.
|
||||
|
||||
Processes schedule operation reports and separates successful operations
|
||||
from failed ones, formatting keys as "namespace/schedule_name" for consistency.
|
||||
|
||||
Args:
|
||||
input_data (list[dict[str, Any]]): The schedule reports. Each item must
|
||||
contain 'namespace', 'schedule_name', 'success', 'message', and optionally
|
||||
'attachment'
|
||||
|
||||
Returns:
|
||||
tuple[list[str], dict[str, Any]]: A tuple of:
|
||||
- Successful schedule keys in the form "namespace/schedule_name"
|
||||
- Error map keyed by the same string to error details
|
||||
"""
|
||||
success_keys = [
|
||||
f'{value["namespace"]}/{value["schedule_name"]}'
|
||||
for value in input_data
|
||||
if value['success']
|
||||
]
|
||||
|
||||
error_keys = {
|
||||
f'{value["namespace"]}/{value["schedule_name"]}': {
|
||||
'message': value['message'],
|
||||
'attachment': value.get('attachment', None),
|
||||
}
|
||||
for value in input_data
|
||||
if not value['success']
|
||||
}
|
||||
|
||||
return success_keys, error_keys
|
||||
|
||||
def parse_report(self, input_data: dict[str, dict[str, Any]]) -> tuple[list[str], list[str]]:
|
||||
"""
|
||||
Parse the report data to extract success and error keys.
|
||||
|
||||
Processes operation reports and separates successful operations from failed ones
|
||||
based on the 'success' field in each report item.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, dict[str, Any]]): The input data containing report items.
|
||||
Each item should have a 'success' field indicating success/failure
|
||||
|
||||
Returns:
|
||||
tuple[list[str], list[str]]: A tuple containing:
|
||||
- List of successful keys
|
||||
- List of error keys
|
||||
"""
|
||||
success_keys = [key for key, value in input_data.items() if value['success']]
|
||||
|
||||
error_keys = [key for key, value in input_data.items() if not value['success']]
|
||||
|
||||
return success_keys, error_keys
|
||||
|
||||
def manage_and_send_report(
|
||||
self,
|
||||
metadata: dict[str, Any],
|
||||
success_keys: list[str],
|
||||
error_keys: dict[str, Any],
|
||||
schedule_type: str,
|
||||
schedule_data: dict[str, Any],
|
||||
):
|
||||
"""
|
||||
Manage and send success and error reports based on the provided keys and data.
|
||||
|
||||
Sends separate notifications for successful and failed operations, formatting
|
||||
error messages with attachments when available.
|
||||
|
||||
Args:
|
||||
metadata (dict[str, Any]): Metadata for logging and notifications
|
||||
success_keys (list[str]): List of keys that were successful
|
||||
error_keys (dict[str, Any]): Dictionary of error keys mapped to error details
|
||||
schedule_type (str): The type of schedule being reported (e.g., 'created schedules')
|
||||
schedule_data (dict[str, Any]): The schedule data containing items and notification ID
|
||||
"""
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
metadata=metadata,
|
||||
message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}',
|
||||
notification_id=schedule_data['id'],
|
||||
attachment=schedule_data['items'],
|
||||
)
|
||||
|
||||
if len(error_keys) > 0:
|
||||
attachment = []
|
||||
for key, value in error_keys.items():
|
||||
if value['attachment'] is not None:
|
||||
attachment.append(f'{key}:\n{value["message"]}\n{value["attachment"]}')
|
||||
else:
|
||||
attachment.append(f'{key}:\n{value["message"]}')
|
||||
|
||||
self.send_error_report(
|
||||
metadata=metadata,
|
||||
message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}',
|
||||
notification_id=f'{schedule_data["id"]}_ERROR',
|
||||
attachment=topic_separator.join(attachment),
|
||||
)
|
||||
|
||||
@activity.defn(name='report_schedule_orchestration')
|
||||
def report_schedule_orchestration(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Report the orchestration result to the notification handler.
|
||||
|
||||
Processes schedule orchestration results and sends notifications for
|
||||
created, updated, and deleted schedules with success and error details.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- created_schedules (list[dict[str, Any]]): The created schedules
|
||||
- updated_schedules (list[dict[str, Any]]): The updated schedules
|
||||
- deleted_schedules (list[dict[str, Any]]): The deleted schedules
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Reporting orchestration...', metadata=metadata)
|
||||
|
||||
created_schedules = input_data['created_schedules']
|
||||
updated_schedules = input_data['updated_schedules']
|
||||
deleted_schedules = input_data['deleted_schedules']
|
||||
|
||||
schedules_report = {
|
||||
'created schedules': {
|
||||
'items': created_schedules,
|
||||
'id': 'REPORT_ORCHESTRATION_CREATED_SCHEDULES',
|
||||
},
|
||||
'updated schedules': {
|
||||
'items': updated_schedules,
|
||||
'id': 'REPORT_ORCHESTRATION_UPDATED_SCHEDULES',
|
||||
},
|
||||
'deleted schedules': {
|
||||
'items': deleted_schedules,
|
||||
'id': 'REPORT_ORCHESTRATION_DELETED_SCHEDULES',
|
||||
},
|
||||
}
|
||||
|
||||
for schedule_type, schedule_data in schedules_report.items():
|
||||
if len(schedule_data['items']) > 0:
|
||||
success_keys, error_keys = self.parse_report_schedule(schedule_data['items'])
|
||||
|
||||
self.manage_and_send_report(
|
||||
metadata=metadata,
|
||||
success_keys=success_keys,
|
||||
error_keys=error_keys,
|
||||
schedule_type=schedule_type,
|
||||
schedule_data=schedule_data,
|
||||
)
|
||||
|
||||
@activity.defn(name='report_slot_orchestration')
|
||||
def report_slot_orchestration(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Report the slot orchestration result to the notification handler.
|
||||
|
||||
Processes slot orchestration results and sends notifications for
|
||||
inserted and deleted slots with success and error details.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- inserted_slots (dict[str, Any]): The inserted slots
|
||||
- deleted_slots (dict[str, Any]): The deleted slots
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Reporting orchestration...', metadata=metadata)
|
||||
|
||||
inserted_slots = input_data['inserted_slots']
|
||||
deleted_slots = input_data['deleted_slots']
|
||||
|
||||
if len(inserted_slots) > 0:
|
||||
success_keys, error_keys = self.parse_report(inserted_slots)
|
||||
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
metadata=metadata,
|
||||
message=f'Inserted slots: \n {", ".join(success_keys)}',
|
||||
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
|
||||
)
|
||||
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
metadata=metadata,
|
||||
message=f'Failed to insert slots: \n {", ".join(error_keys)}',
|
||||
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
|
||||
attachment=inserted_slots,
|
||||
)
|
||||
|
||||
if len(deleted_slots) > 0:
|
||||
success_keys, error_keys = self.parse_report(deleted_slots)
|
||||
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
metadata=metadata,
|
||||
message=f'Deleted slots: \n {", ".join(success_keys)}',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
|
||||
)
|
||||
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
metadata=metadata,
|
||||
message=f'Failed to delete slots: \n {", ".join(error_keys)}',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
|
||||
attachment=deleted_slots,
|
||||
)
|
||||
|
||||
@activity.defn(name='format_log_report')
|
||||
def format_log_report(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Format the receiver_groups status to a dataframe to be stored in the database.
|
||||
|
||||
Transforms receiver group notification data into a structured format suitable
|
||||
for database storage, aggregating notifications by unique notification ID and trigger.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- receiver_groups (dict): The receiver groups configuration with notifications
|
||||
- mail_type (str): The type of mail for the report (Alerts/Reports)
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: The formatted log report as a dictionary representation of a DataFrame
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
mail_type = input_data['mail_type']
|
||||
|
||||
self.info('Formatting log report...', metadata=metadata)
|
||||
|
||||
receiver_groups = input_data['receiver_groups']
|
||||
|
||||
data = {}
|
||||
|
||||
for group_name, group_config in receiver_groups.items():
|
||||
for notification in group_config['notifications']:
|
||||
notification_id = notification['notification_id']
|
||||
trigger = notification['trigger']
|
||||
|
||||
key = f'{notification_id}:{trigger}'
|
||||
|
||||
if key not in data:
|
||||
data[key] = {
|
||||
'status': group_config['status'],
|
||||
'timestamp': notification['timestamp'],
|
||||
'groups': [group_name],
|
||||
'message': notification['message'],
|
||||
'level': notification['level'],
|
||||
'notification_id': notification_id,
|
||||
'block': notification['block'],
|
||||
'schedule': trigger,
|
||||
'pipeline': notification['pipeline'],
|
||||
'project': notification['project'],
|
||||
'model_name': notification['model_name'],
|
||||
'model_id': notification['model_id'],
|
||||
'mail_type': mail_type,
|
||||
}
|
||||
else:
|
||||
if group_name not in data[key]['groups']:
|
||||
data[key]['groups'].append(group_name)
|
||||
|
||||
data_values: DataFrame = DataFrame(list(data.values()))
|
||||
|
||||
# ``DataFrame.to_dict()`` is typed as ``dict[Hashable, Any]`` in pandas
|
||||
# stubs, but default orientation uses column names (str keys).
|
||||
return cast(dict[str, Any], data_values.to_dict())
|
||||
|
||||
@activity.defn(name='filter_notification_reports')
|
||||
def filter_notification_reports(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Filter notifications for comprehensive scheduled reports.
|
||||
|
||||
This method filters notifications of all levels (ERROR, WARNING, INFO, DEBUG)
|
||||
for scheduled report generation. Unlike alert filtering, this method does not
|
||||
implement TTL-based duplicate prevention since reports are meant to provide
|
||||
comprehensive coverage of system activity within a time window.
|
||||
|
||||
Filtering Logic:
|
||||
- Processes all notification levels (not just ERROR)
|
||||
- Applies group-specific content filtering for "reports" type
|
||||
- Respects ignore lists for each receiver group
|
||||
- Prevents duplicate notifications within the same report
|
||||
- Groups notifications by receiver group configurations
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict): Workflow execution metadata for logging
|
||||
- notification_package (list): All notifications to filter (any level)
|
||||
- sending_configs (list): Receiver group configurations with:
|
||||
- group_name (str): Name of the receiver group
|
||||
- contents (list): Content types to include (must contain "reports")
|
||||
- ignore (list, optional): Notification IDs to exclude
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Filtered receiver groups with their notifications, keyed by group_name.
|
||||
Each group contains:
|
||||
- All receiver group configuration fields
|
||||
- notifications (list): Filtered notifications for this group
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
notification_package = input_data['notification_package']
|
||||
sending_configs = input_data['sending_configs']
|
||||
|
||||
self.info('Filtering notification reports...', metadata=metadata)
|
||||
|
||||
receiver_groups = {}
|
||||
|
||||
for receiver_group in sending_configs:
|
||||
group_name = receiver_group['group_name']
|
||||
receiver_groups[group_name] = {**receiver_group, 'notifications': []}
|
||||
receiver_groups[group_name]['notifications'] = []
|
||||
|
||||
already_added_keys = []
|
||||
|
||||
ignore_list = receiver_group.get('ignore', [])
|
||||
|
||||
for notification in notification_package:
|
||||
alert_type = 'reports'
|
||||
notification_id = notification['notification_id']
|
||||
|
||||
key = f'{notification["trigger"]}:{notification_id}'
|
||||
|
||||
# Check if this group must be notified
|
||||
if (
|
||||
alert_type in receiver_group['contents']
|
||||
and notification_id not in ignore_list
|
||||
and key not in already_added_keys
|
||||
):
|
||||
receiver_groups[group_name]['notifications'].append(notification)
|
||||
already_added_keys.append(key)
|
||||
|
||||
# Remove groups with no notifications
|
||||
receiver_groups = {
|
||||
group_name: group
|
||||
for group_name, group in receiver_groups.items()
|
||||
if group['notifications']
|
||||
}
|
||||
|
||||
return receiver_groups
|
||||
511
orchestrator/activities/mongo_db.py
Normal file
511
orchestrator/activities/mongo_db.py
Normal file
@@ -0,0 +1,511 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from datetime import UTC
|
||||
from logging import Logger
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
from sientia_do.temporal.constants import (
|
||||
DATETIME_FORMAT_MS_WITH_TZ,
|
||||
now,
|
||||
)
|
||||
|
||||
|
||||
class MongoDB(SientiaMonitoring):
|
||||
"""
|
||||
MongoDB operations activity for Temporal workflows.
|
||||
|
||||
This class provides MongoDB database operations including document
|
||||
querying, aggregation, timestamp management, and collection management
|
||||
with TTL indexes. It centralizes all MongoDB interactions required by
|
||||
the orchestration system lifecycle.
|
||||
|
||||
Args:
|
||||
connection_string (str): MongoDB connection string
|
||||
database_name (str): Target database name
|
||||
ttl_index_seconds (int): TTL index duration in seconds
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
connection_string: str,
|
||||
database_name: str,
|
||||
ttl_index_seconds: int,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
self.connection_string = connection_string
|
||||
self.database_name = database_name
|
||||
|
||||
self.mongo_db_repository = MongoDBRepository(
|
||||
connection_string=connection_string,
|
||||
database_name=database_name,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
self.ttl_index_seconds = ttl_index_seconds
|
||||
|
||||
# Initialize MongoDB client here (omitted for brevity)
|
||||
logger.info('MongoDB connection initialized')
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Shutdown the MongoDB client and clean up resources.
|
||||
|
||||
Closes the MongoDB repository connection and shuts down the SientiaMonitoring instance.
|
||||
"""
|
||||
self.mongo_db_repository.close()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Ensure the MongoDB client is closed when the object is garbage-collected.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
@activity.defn(
|
||||
name='find_documents_in_mongodb',
|
||||
)
|
||||
def find_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Find documents in a MongoDB collection based on the provided query parameters.
|
||||
|
||||
Args:
|
||||
- input_data (dict): Input data containing query parameters. Contains:
|
||||
- query (dict): Query parameters to filter documents.
|
||||
- timestamp_fields (list[str], optional): Fields to format as RFC3339 with TZ.
|
||||
|
||||
Returns:
|
||||
list[dict]: Documents matching the query with timestamp fields normalized.
|
||||
"""
|
||||
|
||||
query = input_data.get('query', {})
|
||||
metadata = input_data.get('metadata', {})
|
||||
timestamp_fields = input_data.get('timestamp_fields', [])
|
||||
|
||||
collection_name = query.get('collection')
|
||||
if not collection_name:
|
||||
raise ValueError('Collection name must be provided in the query.')
|
||||
|
||||
filters = query.get('filters', {})
|
||||
|
||||
self.info(
|
||||
f"Loading documents from collection '{collection_name}' with filters: {filters}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
try:
|
||||
documents = self.mongo_db_repository.find(collection_name, filters, metadata)
|
||||
|
||||
self.info(
|
||||
f"Loaded {len(documents)} documents from collection '{collection_name}'",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
for document in documents:
|
||||
for timestamp_field in timestamp_fields:
|
||||
if timestamp_field in document:
|
||||
document[timestamp_field] = (
|
||||
document[timestamp_field]
|
||||
.replace(tzinfo=UTC)
|
||||
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
|
||||
)
|
||||
|
||||
self.debug(f'Documents loaded: {documents}', metadata=metadata)
|
||||
|
||||
return documents
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_QUERY_ERROR',
|
||||
message=f'Failed to execute MongoDB query: {e}',
|
||||
block='load_query_from_mongodb',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
raise e
|
||||
|
||||
@activity.defn(name='aggregate_documents_in_mongodb')
|
||||
def aggregate_documents_in_mongodb(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Aggregate documents in a MongoDB collection based on the provided aggregation pipeline.
|
||||
|
||||
Args:
|
||||
- input_data (dict): Input data containing aggregation parameters. Contains:
|
||||
- query (dict): Aggregation parameters including 'collection' and 'aggregation'.
|
||||
- timestamp_fields (list[str], optional): Fields to format as RFC3339 with TZ.
|
||||
|
||||
Returns:
|
||||
list[dict]: Aggregated documents with timestamp fields normalized.
|
||||
"""
|
||||
|
||||
query = input_data.get('query', {})
|
||||
metadata = input_data.get('metadata', {})
|
||||
timestamp_fields = input_data.get('timestamp_fields', [])
|
||||
|
||||
collection_name = query.get('collection')
|
||||
if not collection_name:
|
||||
raise ValueError('Collection name must be provided in the query.')
|
||||
aggregation = query.get('aggregation')
|
||||
if not aggregation:
|
||||
raise ValueError('Aggregation must be provided.')
|
||||
aggregation.append({'$project': {'_id': 0}})
|
||||
|
||||
self.info(
|
||||
f"Aggregating documents from collection '{collection_name}' with aggregation: {aggregation}",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
try:
|
||||
aggregated_documents = self.mongo_db_repository.aggregate(
|
||||
collection_name, aggregation, metadata
|
||||
)
|
||||
|
||||
self.info(
|
||||
f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
for document in aggregated_documents:
|
||||
for timestamp_field in timestamp_fields:
|
||||
if timestamp_field in document:
|
||||
document[timestamp_field] = (
|
||||
document[timestamp_field]
|
||||
.replace(tzinfo=UTC)
|
||||
.strftime(DATETIME_FORMAT_MS_WITH_TZ)
|
||||
)
|
||||
|
||||
self.debug(f'Aggregation result: {aggregated_documents}', metadata=metadata)
|
||||
|
||||
return aggregated_documents
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_AGGREGATION_ERROR',
|
||||
message=f'Failed to execute MongoDB aggregation: {e}',
|
||||
block='aggregate_documents_in_mongodb',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
raise e
|
||||
|
||||
@activity.defn(name='update_pipelines_timestamps')
|
||||
def update_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Update `updated_at` timestamps for successfully updated pipelines.
|
||||
|
||||
Updates the `updated_at` field in the `orchestrated_schedules` collection
|
||||
for all pipelines that were successfully updated in Temporal.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- updated_pipelines (list[dict]): Pipelines with success flags to consider
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
updated_pipelines = input_data.get('updated_pipelines', [])
|
||||
metadata = input_data.get('metadata', {})
|
||||
date_now = now()
|
||||
|
||||
self.info('Updating pipelines timestamps...', metadata=metadata)
|
||||
|
||||
argument = [
|
||||
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
|
||||
for pipeline in updated_pipelines
|
||||
if pipeline['success']
|
||||
]
|
||||
data_filter = {'$or': argument} if argument else {}
|
||||
|
||||
try:
|
||||
self.mongo_db_repository.update_many(
|
||||
'orchestrated_schedules', data_filter, {'$set': {'updated_at': date_now}}, metadata
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
|
||||
message=f'Failed to update pipelines timestamps: {e}',
|
||||
block='update_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(
|
||||
f'Updated {len(updated_pipelines)} pipelines timestamps',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@activity.defn(name='create_pipelines_timestamps')
|
||||
def create_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Insert `updated_at` timestamps for newly created pipelines.
|
||||
|
||||
Inserts new documents into the `orchestrated_schedules` collection
|
||||
for all pipelines that were successfully created in Temporal.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- created_pipelines (list[dict]): Pipelines with success flags to consider
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
created_pipelines = input_data.get('created_pipelines', [])
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Creating pipelines timestamps...', metadata=metadata)
|
||||
|
||||
date_now = now()
|
||||
|
||||
argument = [
|
||||
{
|
||||
'schedule_name': pipeline['schedule_name'],
|
||||
'namespace': pipeline['namespace'],
|
||||
'updated_at': date_now,
|
||||
}
|
||||
for pipeline in created_pipelines
|
||||
if pipeline['success']
|
||||
]
|
||||
data_filter = argument if argument else {}
|
||||
|
||||
try:
|
||||
if data_filter:
|
||||
self.mongo_db_repository.insert_many(
|
||||
'orchestrated_schedules', data_filter, metadata
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_CREATE_PIPELINES_ERROR',
|
||||
message=f'Failed to create pipelines timestamps: {e}',
|
||||
block='create_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(
|
||||
f'Created {len(created_pipelines)} pipelines timestamps',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@activity.defn(name='delete_pipelines_timestamps')
|
||||
def delete_pipelines_timestamps(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Delete timestamp rows for successfully deleted pipelines.
|
||||
|
||||
Removes documents from the `orchestrated_schedules` collection
|
||||
for all pipelines that were successfully deleted from Temporal.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- deleted_pipelines (list[dict]): Pipelines with success flags to consider
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
deleted_pipelines = input_data.get('deleted_pipelines', [])
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Deleting pipelines timestamps...', metadata=metadata)
|
||||
|
||||
argument = [
|
||||
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
|
||||
for pipeline in deleted_pipelines
|
||||
if pipeline['success']
|
||||
]
|
||||
data_filter = {'$or': argument} if argument else {}
|
||||
|
||||
try:
|
||||
self.mongo_db_repository.delete_many('orchestrated_schedules', data_filter, metadata)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_DELETE_PIPELINES_ERROR',
|
||||
message=f'Failed to delete pipelines timestamps: {e}',
|
||||
block='delete_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(
|
||||
f'Deleted {len(deleted_pipelines)} pipelines timestamps',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
@activity.defn(name='create_collection_with_ttl_index')
|
||||
def create_collection_with_ttl_index(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Create collections with TTL indexes for pipeline topics.
|
||||
|
||||
Creates MongoDB collections for scouter pipeline topics and sets up
|
||||
TTL indexes on the `inserted_at` field to automatically expire old documents.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- pipelines (dict[str, Any]): Pipeline configurations with topic names
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
"""
|
||||
pipelines = input_data.get('pipelines', {})
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info(
|
||||
f'Creating collection with TTL index for pipelines: {list(pipelines.keys())}',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
collection_names = self.mongo_db_repository.database.list_collection_names()
|
||||
|
||||
created_collections = []
|
||||
created_indexes = []
|
||||
|
||||
for _pipeline_name, pipeline_config in pipelines.items():
|
||||
collection = pipeline_config.get('topic', None)
|
||||
if not collection:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Check if collection exists
|
||||
if collection not in collection_names:
|
||||
self.mongo_db_repository.database.create_collection(collection)
|
||||
created_collections.append(collection)
|
||||
|
||||
collection = self.mongo_db_repository.database[collection]
|
||||
# Check if TTL index exists
|
||||
existing_indexes = collection.list_indexes()
|
||||
ttl_index_exists = False
|
||||
for index in existing_indexes:
|
||||
if (
|
||||
'inserted_at' in index['key']
|
||||
and index.get('expireAfterSeconds') is not None
|
||||
):
|
||||
ttl_index_exists = True
|
||||
break
|
||||
|
||||
# Create TTL index if it doesn't exist
|
||||
if not ttl_index_exists:
|
||||
collection.create_index(
|
||||
'inserted_at', expireAfterSeconds=self.ttl_index_seconds, background=True
|
||||
)
|
||||
created_indexes.append(collection)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGODB_CREATE_COLLECTION_ERROR',
|
||||
message=f'Failed to create collection {collection} with TTL index: {e}',
|
||||
block='create_collection_with_ttl_index',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(
|
||||
f'Created {len(created_collections)} collections and {len(created_indexes)} indexes',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.debug(f'Created collections: {created_collections}', metadata=metadata)
|
||||
self.debug(f'Created indexes: {created_indexes}', metadata=metadata)
|
||||
|
||||
@activity.defn(name='load_latest_data')
|
||||
def load_latest_data(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Load the latest data from MongoDB collection since a specified timestamp.
|
||||
|
||||
This activity retrieves data from a MongoDB collection, optionally
|
||||
filtering by timestamp to enable incremental data processing. It
|
||||
handles connection management and provides comprehensive error reporting.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Activity input parameters.
|
||||
Required fields:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- collection_name (str): Name of the MongoDB collection
|
||||
- last_data_timestamp (str | None): Last processed timestamp for filtering
|
||||
- base_data_filter (dict[str, Any]): Base query filter conditions
|
||||
|
||||
Returns:
|
||||
list[dict[str, Any]]: Retrieved data, or empty list if no data found
|
||||
|
||||
Raises:
|
||||
Exception: If MongoDB operation fails
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
collection_name = input_data['collection_name']
|
||||
last_data_timestamp = input_data['last_data_timestamp']
|
||||
base_data_filter = input_data['base_data_filter']
|
||||
|
||||
self.debug(f'Loading data from MongoDB: {input_data}', metadata=metadata)
|
||||
|
||||
try:
|
||||
if last_data_timestamp is None:
|
||||
data_filter = base_data_filter
|
||||
else:
|
||||
# ``notification_queue.timestamp`` is stored as a string in
|
||||
# ``DATETIME_FORMAT_WITH_TZ`` (``Notification`` writes it as
|
||||
# ``now().strftime(DATETIME_FORMAT_WITH_TZ)``). Coercing
|
||||
# ``last_data_timestamp`` to ``datetime`` here would force a
|
||||
# BSON ``String`` vs ``Date`` comparison, which always yields
|
||||
# ``False`` (``String < Date`` in BSON sort order) and breaks
|
||||
# incremental loading entirely. Comparing strings preserves the
|
||||
# intended chronological filter because the format is
|
||||
# lexicographically ordered when the timezone is fixed
|
||||
# (``Notification.timestamp`` always uses UTC).
|
||||
data_filter = {
|
||||
**base_data_filter,
|
||||
'timestamp': {'$gt': last_data_timestamp},
|
||||
}
|
||||
|
||||
self.debug(f'Data filter: {data_filter}', metadata=metadata)
|
||||
|
||||
data = self.mongo_db_repository.find(collection_name, data_filter, metadata)
|
||||
|
||||
self.debug(f'Collected: {data}', metadata=metadata)
|
||||
|
||||
self.info(f'Loaded {len(data)} documents from MongoDB', metadata=metadata)
|
||||
|
||||
self.debug(f'Loaded data: {data}', metadata=metadata)
|
||||
|
||||
return data
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MONGO_LOAD_ERROR',
|
||||
message=f'Error loading data from MongoDB: {e}',
|
||||
block='load_latest_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
raise e
|
||||
468
orchestrator/activities/slot_manager.py
Normal file
468
orchestrator/activities/slot_manager.py
Normal file
@@ -0,0 +1,468 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import json
|
||||
import traceback
|
||||
from datetime import datetime, timedelta
|
||||
from logging import Logger
|
||||
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.observability.sientia_monitoring import MetricsController, SientiaMonitoring
|
||||
from sientia_do.repository.redis_repository_sync import RedisRepository
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
|
||||
|
||||
|
||||
class SlotManager(SientiaMonitoring):
|
||||
"""
|
||||
Redis-based OPC slot management and notification filtering activity.
|
||||
|
||||
This class manages OPC server slots and provides advanced notification
|
||||
filtering capabilities through Redis operations. It handles OPC slot
|
||||
lifecycle management, active ingestor tracking, and implements intelligent
|
||||
notification filtering with TTL-based duplicate prevention.
|
||||
|
||||
Key Features:
|
||||
- OPC slot loading, updating, and deletion
|
||||
- Active ingestor management
|
||||
- Notification filtering for alerts with TTL management
|
||||
- Persistent alert detection for ongoing issues
|
||||
- Notification caching with configurable expiration
|
||||
- Group-based filtering with ignore list support
|
||||
|
||||
Args:
|
||||
host (str): Redis server hostname
|
||||
port (int): Redis server port number
|
||||
username (str): Redis authentication username
|
||||
password (str): Redis authentication password
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
username: str,
|
||||
password: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
self.redis_repository = RedisRepository(
|
||||
host=host,
|
||||
port=port,
|
||||
username=username,
|
||||
password=password,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Close the SlotManager connection and clean up resources.
|
||||
|
||||
Closes the Redis repository connection and shuts down the SientiaMonitoring instance.
|
||||
"""
|
||||
self.redis_repository.close()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Ensure the SlotManager connection is closed when the object is garbage-collected.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='load_opc_slots')
|
||||
def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Load all OPC slots from Redis for current system state assessment.
|
||||
|
||||
This method retrieves all OPC server slot configurations from Redis,
|
||||
which are used to determine current resource allocation and identify
|
||||
changes needed for pipeline orchestration.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Activity input containing metadata
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Dictionary of OPC slots keyed by server ID, where each slot contains:
|
||||
- Configuration parameters for OPC server connections
|
||||
- Active pipeline assignments
|
||||
- Resource allocation details
|
||||
|
||||
Raises:
|
||||
Exception: If Redis connection fails or data retrieval errors occur
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Loading OPC slots...', metadata=metadata)
|
||||
|
||||
opc_slots = {}
|
||||
|
||||
try:
|
||||
slot_keys = self.redis_repository.keys('slot:opc_tags:*')
|
||||
|
||||
self.debug(f'Slot keys: {slot_keys}', metadata=metadata)
|
||||
|
||||
if slot_keys:
|
||||
if isinstance(slot_keys[0], bytes):
|
||||
decoded_keys = [key.decode('utf-8') for key in slot_keys]
|
||||
else:
|
||||
decoded_keys = slot_keys
|
||||
|
||||
for key in decoded_keys:
|
||||
opc_slots[key] = self.redis_repository.get(key)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message=f'Failed to load OPC slots: {e}',
|
||||
block='load_opc_slots',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(f'Loaded {len(opc_slots)} OPC slots', metadata=metadata)
|
||||
|
||||
return opc_slots
|
||||
|
||||
@activity.defn(name='load_active_ingestors')
|
||||
def load_active_ingestors(self, input_data: dict[str, Any]) -> list[str]:
|
||||
"""
|
||||
Load all active ingestors from Redis.
|
||||
|
||||
Retrieves all active ingestor heartbeat keys from Redis to determine
|
||||
which ingestors are currently available for slot assignment.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Activity input containing metadata
|
||||
|
||||
Returns:
|
||||
list[str]: A list of active ingestor keys from Redis
|
||||
"""
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Loading active ingestors...', metadata=metadata)
|
||||
|
||||
try:
|
||||
active_ingestors = self.redis_repository.keys('heartbeat:ingestor:*')
|
||||
|
||||
self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata)
|
||||
|
||||
self.debug(f'Active ingestors: \n {active_ingestors}', metadata=metadata)
|
||||
|
||||
ingestors = []
|
||||
|
||||
for ingestor in active_ingestors:
|
||||
if isinstance(ingestor, bytes):
|
||||
ingestors.append(ingestor.decode('utf-8'))
|
||||
else:
|
||||
ingestors.append(ingestor)
|
||||
|
||||
return ingestors
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message=f'Failed to load active ingestors: {e}',
|
||||
block='load_active_ingestors',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
@activity.defn(name='update_slots')
|
||||
def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Update OPC slots in Redis
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the slots to update.
|
||||
- to_insert (dict[str, Any]): The slots to insert.
|
||||
|
||||
Returns:
|
||||
- report (dict[str, Any]): A report of the updated slots.
|
||||
"""
|
||||
|
||||
to_insert = input_data['to_insert']
|
||||
metadata = input_data.get('metadata', {})
|
||||
self.info('Updating OPC slots...', metadata=metadata)
|
||||
|
||||
report = {}
|
||||
|
||||
success_count = 0
|
||||
|
||||
for slot in to_insert:
|
||||
try:
|
||||
self.redis_repository.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
|
||||
report[slot] = {'success': True, 'message': 'Slot updated successfully'}
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
self.error(f'Failed to update slot {slot}: {str(e)}', metadata=metadata)
|
||||
report[slot] = {'success': False, 'message': str(e)}
|
||||
|
||||
self.info(f'Updated {success_count} of {len(to_insert)} OPC slots', metadata=metadata)
|
||||
|
||||
self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name='delete_slots')
|
||||
def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Delete OPC slots from Redis
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the slots to delete.
|
||||
- to_delete (list[str]): The slots to delete.
|
||||
|
||||
Returns:
|
||||
- report (dict[str, Any]): A report of the deleted slots.
|
||||
"""
|
||||
|
||||
to_delete = input_data['to_delete']
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
self.info('Deleting OPC slots...', metadata=metadata)
|
||||
|
||||
report = {}
|
||||
|
||||
success_count = 0
|
||||
|
||||
for slot in to_delete:
|
||||
try:
|
||||
self.redis_repository.delete(f'slot:opc_tags:{slot}')
|
||||
report[slot] = {'success': True, 'message': 'Slot deleted successfully'}
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
self.error(f'Failed to delete slot {slot}: {str(e)}', metadata=metadata)
|
||||
report[slot] = {'success': False, 'message': str(e)}
|
||||
|
||||
self.info(f'Deleted {success_count} of {len(to_delete)} OPC slots', metadata=metadata)
|
||||
|
||||
self.debug(f'Report: \n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name='get_last_data_timestamp')
|
||||
def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Get the last data timestamp from Redis.
|
||||
|
||||
Retrieves the last processed timestamp for a specific mail type from Redis.
|
||||
This timestamp is used for incremental data loading to avoid reprocessing
|
||||
already processed notifications.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
- mail_type (str): The type of mail to get timestamp for (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
str | None: The last data timestamp as a string, or None if no timestamp exists
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
key = f'notification_last_timestamp:{input_data["mail_type"]}'
|
||||
|
||||
try:
|
||||
data_hold = self.redis_repository.get(key)
|
||||
except Exception as e:
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message=f'Error getting last data timestamp: {e}',
|
||||
block='get_last_data_timestamp',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc(),
|
||||
)
|
||||
raise e
|
||||
|
||||
self.debug(f'Last collected timestamp: {data_hold}', metadata=metadata)
|
||||
|
||||
if not data_hold:
|
||||
return None
|
||||
|
||||
return data_hold
|
||||
|
||||
@activity.defn(name='put_last_data_timestamp')
|
||||
def put_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Store the last data timestamp in Redis.
|
||||
|
||||
Extracts the maximum timestamp from the provided data and stores it in Redis
|
||||
with a TTL for the specified mail type. This enables incremental processing
|
||||
in subsequent workflow executions.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
- data (list[dict]): The data to extract timestamp from
|
||||
- mail_type (str): The type of mail to store timestamp for (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
str | None: The last data timestamp that was stored, or None if no data exists
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
key = f'notification_last_timestamp:{input_data["mail_type"]}'
|
||||
|
||||
data = DataFrame(input_data['data'])
|
||||
|
||||
if data.empty:
|
||||
self.warning('No data to insert', metadata=metadata)
|
||||
return None
|
||||
|
||||
last_data_timestamp = data['timestamp'].max()
|
||||
|
||||
self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
|
||||
|
||||
try:
|
||||
self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
|
||||
except Exception as e:
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='REDIS_SET_ERROR',
|
||||
message=f'Error setting last data timestamp: {e}',
|
||||
block='put_last_data_timestamp',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc(),
|
||||
)
|
||||
raise e
|
||||
|
||||
return last_data_timestamp
|
||||
|
||||
@activity.defn(name='filter_notification_alerts')
|
||||
def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Filter notification alerts with intelligent TTL-based duplicate prevention.
|
||||
|
||||
This method implements advanced notification filtering for ERROR-level alerts,
|
||||
preventing spam through TTL management and detecting persistent issues that
|
||||
require escalation. It applies user group-based filtering with configurable
|
||||
ignore lists and content policies.
|
||||
|
||||
Filtering Logic:
|
||||
- Checks Redis cache for recently sent notifications
|
||||
- Identifies "core_alerts" for new notifications
|
||||
- Detects "persistent_alerts" for ongoing issues beyond TTL
|
||||
- Applies group-specific content filtering and ignore lists
|
||||
- Prevents duplicate notifications within the same TTL window
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict): Workflow execution metadata for logging
|
||||
- notification_package (list): ERROR-level notifications to filter
|
||||
- sending_configs (list): Receiver group configurations with:
|
||||
- group_name (str): Name of the receiver group
|
||||
- contents (list): Alert types to include (core_alerts, persistent_alerts)
|
||||
- ignore (list, optional): Notification IDs to exclude
|
||||
- notification_ttl (int): Seconds before considering notification persistent
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Filtered receiver groups with their notifications, keyed by group_name.
|
||||
Each group contains:
|
||||
- All receiver group configuration fields
|
||||
- notifications (list): Filtered notifications for this group
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
notification_package = input_data['notification_package']
|
||||
sending_configs = input_data['sending_configs']
|
||||
notification_ttl = input_data['notification_ttl']
|
||||
|
||||
self.info('Filtering notification alerts...', metadata=metadata)
|
||||
|
||||
receiver_groups = {}
|
||||
|
||||
for receiver_group in sending_configs:
|
||||
group_name = receiver_group['group_name']
|
||||
receiver_groups[group_name] = {**receiver_group, 'notifications': []}
|
||||
receiver_groups[group_name]['notifications'] = []
|
||||
|
||||
already_added_keys = []
|
||||
|
||||
ignore_list = receiver_group.get('ignore', [])
|
||||
|
||||
for notification in notification_package:
|
||||
alert_type = 'do_nothing'
|
||||
notification_id = notification['notification_id']
|
||||
# Check if notification was recently sent
|
||||
key = f'{notification["trigger"]}:{notification_id}'
|
||||
|
||||
last_sent = self.redis_repository.get(key)
|
||||
|
||||
if last_sent is None:
|
||||
alert_type = 'core_alerts'
|
||||
|
||||
else:
|
||||
last_sent = datetime.strptime(last_sent, DATETIME_FORMAT_MS_WITH_TZ)
|
||||
|
||||
# Check if "notification_ttl" seconds has passed since last sent
|
||||
if (now() - last_sent) > timedelta(seconds=notification_ttl):
|
||||
alert_type = 'persistent_alerts'
|
||||
|
||||
# Check if this group must be notified
|
||||
if (
|
||||
alert_type in receiver_group['contents']
|
||||
and notification_id not in ignore_list
|
||||
and key not in already_added_keys
|
||||
):
|
||||
receiver_groups[group_name]['notifications'].append(notification)
|
||||
already_added_keys.append(key)
|
||||
|
||||
# Remove groups with no notifications
|
||||
receiver_groups = {
|
||||
group_name: group
|
||||
for group_name, group in receiver_groups.items()
|
||||
if group['notifications']
|
||||
}
|
||||
|
||||
return receiver_groups
|
||||
|
||||
@activity.defn(name='store_notification_cache')
|
||||
def store_notification_cache(self, input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Store notification cache in Redis to track recently sent notifications.
|
||||
|
||||
Stores successfully sent notifications in Redis with a TTL to prevent
|
||||
duplicate alert delivery. Only notifications with status 'sent' are cached.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict): Metadata for logging purposes
|
||||
- log_report (list[dict]): The log report containing notification statuses
|
||||
- sent_ttl (int): Time to live for sent notification cache in seconds
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
log_report = DataFrame(input_data['log_report'])
|
||||
sent_ttl = input_data['sent_ttl']
|
||||
|
||||
self.info('Storing notification cache...', metadata=metadata)
|
||||
|
||||
date_now = now().strftime(DATETIME_FORMAT_MS_WITH_TZ)
|
||||
|
||||
for _index, row in log_report.iterrows():
|
||||
status = row['status']
|
||||
if status == 'sent':
|
||||
key = f'{row["schedule"]}:{row["notification_id"]}'
|
||||
self.redis_repository.set(key, date_now, ttl=sent_ttl)
|
||||
|
||||
self.info('Notification cache stored...', metadata=metadata)
|
||||
484
orchestrator/activities/temporal_manager.py
Normal file
484
orchestrator/activities/temporal_manager.py
Normal file
@@ -0,0 +1,484 @@
|
||||
from temporalio import activity, workflow
|
||||
from temporalio.client import (
|
||||
Client,
|
||||
Schedule,
|
||||
ScheduleActionStartWorkflow,
|
||||
ScheduleIntervalSpec,
|
||||
ScheduleSpec,
|
||||
ScheduleUpdate,
|
||||
ScheduleUpdateInput,
|
||||
)
|
||||
from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import json
|
||||
import traceback
|
||||
from datetime import timedelta
|
||||
from logging import Logger
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
|
||||
from sientia_do.temporal.worker.prepare_worker import build_queue_name
|
||||
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
|
||||
RUNTIME_WORKFLOWS = ['predictions_batch', 'minimal_retrain', 'drift', 'simple_metrics']
|
||||
|
||||
|
||||
class TemporalManager(SientiaMonitoring):
|
||||
"""
|
||||
Temporal workflow and schedule management activity.
|
||||
|
||||
This class manages Temporal schedules across multiple namespaces,
|
||||
providing operations for schedule creation, updates, deletion, and
|
||||
normalization. It handles connections to both scouter and laborious
|
||||
namespaces for comprehensive workflow orchestration.
|
||||
|
||||
Args:
|
||||
host (str): Temporal server host address
|
||||
scouter_namespace (str): Scouter workflow namespace
|
||||
laborious_namespace (str): Laborious workflow namespace
|
||||
logger (Logger): Application logger instance
|
||||
notification_handler (NotificationHandler): Notification management handler
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
scouter_namespace: str,
|
||||
laborious_namespace: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
self.temporal_host = host
|
||||
self.scouter_namespace = scouter_namespace
|
||||
self.laborious_namespace = laborious_namespace
|
||||
self.temporal_clients: dict[str, Client] = {}
|
||||
|
||||
self.model_id_id_key = SearchAttributeKey.for_keyword('model_id')
|
||||
self.model_name_id_key = SearchAttributeKey.for_keyword('model_name')
|
||||
self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated')
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Close the TemporalManager connection and clean up resources.
|
||||
|
||||
Shuts down the SientiaMonitoring instance and releases all resources.
|
||||
"""
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Ensure the TemporalManager connection is closed when the object is garbage-collected.
|
||||
"""
|
||||
self.close()
|
||||
|
||||
async def connect_to_temporal(self):
|
||||
"""
|
||||
Connect to Temporal server namespaces used by scouter and laborious workflows.
|
||||
|
||||
Creates and caches `Client` connections for both namespaces for later use.
|
||||
The connections are stored in `temporal_clients` dictionary for efficient access.
|
||||
"""
|
||||
self.logger.info(f'Connecting to Temporal side namespaces at {self.temporal_host}')
|
||||
self.logger.info(f'Scouter namespace: {self.scouter_namespace}')
|
||||
|
||||
scouter_client = await Client.connect(
|
||||
target_host=self.temporal_host, namespace=self.scouter_namespace
|
||||
)
|
||||
|
||||
self.logger.info(f'Laborious namespace: {self.laborious_namespace}')
|
||||
|
||||
laborious_client = await Client.connect(
|
||||
target_host=self.temporal_host, namespace=self.laborious_namespace
|
||||
)
|
||||
|
||||
self.temporal_clients = {
|
||||
self.scouter_namespace: scouter_client,
|
||||
self.laborious_namespace: laborious_client,
|
||||
}
|
||||
|
||||
@activity.defn(name='normalize_schedules')
|
||||
async def normalize_schedules(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Normalize schedules by removing orphaned schedules from Temporal.
|
||||
|
||||
Any schedule marked with search attribute `orchestrated=true` that does not
|
||||
exist in MongoDB collection `orchestrated_schedules` will be deleted.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- orchestrated_schedules (dict[str, Any]): Current orchestrated schedules from MongoDB.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
remove_count = 0
|
||||
|
||||
self.info('Getting orchestrated schedules...', metadata=metadata)
|
||||
|
||||
orchestrated_schedules = input_data.get('orchestrated_schedules', {})
|
||||
|
||||
for namespace, client in self.temporal_clients.items():
|
||||
try:
|
||||
schedules = orchestrated_schedules.get(namespace, {})
|
||||
|
||||
self.info(f'Getting orchestrated schedules for {namespace}', metadata=metadata)
|
||||
|
||||
async for schedule in await client.list_schedules():
|
||||
search_attrs = getattr(schedule, 'search_attributes', {})
|
||||
if search_attrs.get('orchestrated', ['false']) == ['true']:
|
||||
schedule_id = schedule.id
|
||||
|
||||
if schedule_id not in schedules:
|
||||
self.info(
|
||||
f'Schedule {schedule_id} not found in mongo db, cleaning up',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
handle = client.get_schedule_handle(schedule_id)
|
||||
|
||||
await handle.delete()
|
||||
|
||||
remove_count += 1
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR',
|
||||
message=f'Failed to normalize schedules: {e}',
|
||||
block='normalize_schedules',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
self.info(f'Removed {remove_count} schedules', metadata=metadata)
|
||||
|
||||
@activity.defn(name='create_schedules')
|
||||
async def create_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Create schedules in Temporal.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to create.
|
||||
- schedules (dict[str, Any]): The schedules to create.
|
||||
|
||||
Returns:
|
||||
- list[dict[str, Any]]: Report entries for each attempted schedule creation.
|
||||
"""
|
||||
|
||||
schedules_to_create = input_data['schedules']
|
||||
metadata = input_data.get('metadata', {})
|
||||
|
||||
report = []
|
||||
|
||||
success_count = 0
|
||||
|
||||
self.info('Creating schedules...', metadata=metadata)
|
||||
|
||||
for namespace, schedules in schedules_to_create.items():
|
||||
client = self.temporal_clients.get(namespace)
|
||||
|
||||
if not client:
|
||||
raise ValueError(
|
||||
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
|
||||
)
|
||||
|
||||
for schedule_name, schedule in schedules.items():
|
||||
search_attributes = TypedSearchAttributes(
|
||||
[
|
||||
SearchAttributePair(key=self.model_id_id_key, value=schedule['model_id']),
|
||||
SearchAttributePair(
|
||||
key=self.model_name_id_key, value=schedule['model_name']
|
||||
),
|
||||
SearchAttributePair(key=self.orchestrated_id_key, value='true'),
|
||||
]
|
||||
)
|
||||
workflow_type = schedule['workflow_type']
|
||||
|
||||
try:
|
||||
execution_timeout_seconds = schedule.get('execution_timeout_seconds', 300)
|
||||
task_timeout_seconds = schedule.get('task_timeout_seconds', 300)
|
||||
self.debug(f'Creating schedule {schedule_name}:', metadata=metadata)
|
||||
self.debug(
|
||||
f'{json.dumps(schedule, indent=4, sort_keys=True)}', metadata=metadata
|
||||
)
|
||||
|
||||
task_queue_name = self._build_task_queue_name(workflow_type, schedule)
|
||||
schedule['task_queue'] = task_queue_name
|
||||
|
||||
await client.create_schedule(
|
||||
schedule_name,
|
||||
Schedule(
|
||||
action=ScheduleActionStartWorkflow(
|
||||
workflow_type,
|
||||
schedule,
|
||||
id=schedule_name,
|
||||
task_queue=task_queue_name,
|
||||
execution_timeout=timedelta(seconds=execution_timeout_seconds),
|
||||
run_timeout=timedelta(seconds=execution_timeout_seconds),
|
||||
task_timeout=timedelta(seconds=task_timeout_seconds),
|
||||
typed_search_attributes=search_attributes,
|
||||
),
|
||||
spec=ScheduleSpec(
|
||||
intervals=[
|
||||
ScheduleIntervalSpec(
|
||||
every=timedelta(
|
||||
seconds=parse_frequency(schedule.get('frequency', '1m'))
|
||||
),
|
||||
offset=timedelta(
|
||||
seconds=parse_frequency(schedule.get('offset', '0m'))
|
||||
),
|
||||
)
|
||||
]
|
||||
),
|
||||
),
|
||||
search_attributes=search_attributes,
|
||||
)
|
||||
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': True,
|
||||
'message': 'Schedule created successfully',
|
||||
}
|
||||
)
|
||||
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
self.error(
|
||||
f'Failed to create schedule {schedule_name}: {str(e)}', metadata=metadata
|
||||
)
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': False,
|
||||
'message': str(e),
|
||||
}
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Created {success_count} of {len(schedules_to_create)} schedules', metadata=metadata
|
||||
)
|
||||
|
||||
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
|
||||
|
||||
return report
|
||||
|
||||
@staticmethod
|
||||
def _build_task_queue_name(workflow_type: str, schedule: dict[str, Any]) -> str:
|
||||
"""
|
||||
Build the Temporal task queue name for a schedule.
|
||||
|
||||
Only workflow types listed in ``RUNTIME_WORKFLOWS`` get an environment/tenant
|
||||
specific ``runtime`` suffix; every other workflow type gets a plain queue name.
|
||||
"""
|
||||
runtime_name = (
|
||||
schedule.get('runtime', 'legacy') if workflow_type in RUNTIME_WORKFLOWS else None
|
||||
)
|
||||
return build_queue_name(workflow_type, runtime_name)
|
||||
|
||||
def _make_schedule_updater(self, schedule: dict[str, Any], metadata: dict[str, Any]):
|
||||
"""Build the ``ScheduleUpdate`` callback used by ``handler.update`` for one schedule."""
|
||||
|
||||
# fmt: off
|
||||
async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR
|
||||
schedule_action = input_data.description.schedule.action
|
||||
|
||||
self.debug("Updating schedule:", metadata=metadata)
|
||||
|
||||
if hasattr(schedule_action, "args"):
|
||||
self.debug("New schedule:", metadata=metadata)
|
||||
self.debug(
|
||||
f"{json.dumps(schedule, indent=4, sort_keys=True)}", metadata=metadata) # NOSONAR
|
||||
|
||||
schedule_action.args = [schedule]
|
||||
|
||||
input_data.description.schedule.spec.intervals = [
|
||||
ScheduleIntervalSpec(
|
||||
every=timedelta(
|
||||
seconds=parse_frequency(schedule.get('frequency', '1m'))),
|
||||
offset=timedelta(
|
||||
seconds=parse_frequency(schedule.get('offset', '0m'))),
|
||||
)
|
||||
]
|
||||
|
||||
return ScheduleUpdate(schedule=input_data.description.schedule)
|
||||
|
||||
# fmt: on
|
||||
return update_schedule
|
||||
|
||||
async def _update_single_schedule(
|
||||
self,
|
||||
client: Client,
|
||||
schedule_name: str,
|
||||
schedule: dict[str, Any],
|
||||
metadata: dict[str, Any],
|
||||
) -> None:
|
||||
"""Update a single schedule in Temporal, raising if the schedule handle is missing."""
|
||||
handler = client.get_schedule_handle(schedule_name)
|
||||
|
||||
if not handler:
|
||||
raise ValueError(f'Schedule {schedule_name} not found')
|
||||
|
||||
workflow_type = schedule['workflow_type']
|
||||
schedule['task_queue'] = self._build_task_queue_name(workflow_type, schedule)
|
||||
|
||||
update_schedule = self._make_schedule_updater(schedule, metadata)
|
||||
await handler.update(update_schedule)
|
||||
|
||||
@activity.defn(name='update_schedules')
|
||||
async def update_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Update schedules in Temporal.
|
||||
|
||||
Does not modify ``task_queue``. The ``ScheduleUpdate`` callback only patches
|
||||
workflow ``args`` and schedule ``intervals``. A ``runtime`` change requires
|
||||
delete-then-create on the next orchestrator tick.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to update.
|
||||
- schedules (dict[str, Any]): The schedules to update.
|
||||
|
||||
Returns:
|
||||
- list[dict[str, Any]]: Report entries for each attempted schedule update.
|
||||
"""
|
||||
|
||||
schedules_to_update = input_data['schedules']
|
||||
metadata = input_data.get('metadata', {})
|
||||
report = []
|
||||
|
||||
success_count = 0
|
||||
|
||||
self.info('Updating schedules...', metadata=metadata)
|
||||
|
||||
for namespace, schedules in schedules_to_update.items():
|
||||
client = self.temporal_clients.get(namespace)
|
||||
|
||||
if not client:
|
||||
raise ValueError(
|
||||
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
|
||||
)
|
||||
|
||||
for schedule_name, schedule in schedules.items():
|
||||
try:
|
||||
await self._update_single_schedule(client, schedule_name, schedule, metadata)
|
||||
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': True,
|
||||
'message': 'Schedule updated successfully',
|
||||
}
|
||||
)
|
||||
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
self.error(
|
||||
f'Failed to update schedule {schedule_name}: {str(e)}', metadata=metadata
|
||||
)
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': False,
|
||||
'message': str(e),
|
||||
}
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Updated {success_count} of {len(schedules_to_update)} schedules', metadata=metadata
|
||||
)
|
||||
|
||||
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name='delete_schedules')
|
||||
async def delete_schedules(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Delete schedules in Temporal.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to delete.
|
||||
- schedules (list[str]): The schedules to delete.
|
||||
|
||||
Returns:
|
||||
- list[dict[str, Any]]: Report entries for each attempted schedule deletion.
|
||||
"""
|
||||
|
||||
schedules_to_delete = input_data['schedules']
|
||||
metadata = input_data.get('metadata', {})
|
||||
report = []
|
||||
|
||||
success_count = 0
|
||||
|
||||
self.info('Deleting schedules...', metadata=metadata)
|
||||
|
||||
for namespace, schedules in schedules_to_delete.items():
|
||||
client = self.temporal_clients.get(namespace)
|
||||
|
||||
if not client:
|
||||
raise ValueError(
|
||||
f'Temporal client for {namespace} not found, clients: {self.temporal_clients}'
|
||||
)
|
||||
|
||||
for schedule_name in schedules:
|
||||
try:
|
||||
handler = client.get_schedule_handle(schedule_name)
|
||||
|
||||
if not handler:
|
||||
raise ValueError(f'Schedule {schedule_name} not found')
|
||||
|
||||
await handler.delete()
|
||||
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': True,
|
||||
'message': 'Schedule deleted successfully',
|
||||
}
|
||||
)
|
||||
|
||||
success_count += 1
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.error(
|
||||
f'Failed to delete schedule {schedule_name}: {str(e)}', metadata=metadata
|
||||
)
|
||||
report.append(
|
||||
{
|
||||
'namespace': namespace,
|
||||
'schedule_name': schedule_name,
|
||||
'success': False,
|
||||
'message': str(e),
|
||||
'attachment': trace,
|
||||
}
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Deleted {success_count} of {len(schedules_to_delete)} schedules', metadata=metadata
|
||||
)
|
||||
|
||||
self.debug(f'\n {json.dumps(report, indent=4, sort_keys=True)}', metadata=metadata)
|
||||
|
||||
return report
|
||||
23
orchestrator/metrics.py
Normal file
23
orchestrator/metrics.py
Normal file
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
Prometheus metric definitions for the orchestrator application.
|
||||
|
||||
This module exposes Prometheus counters and gauges for monitoring the
|
||||
orchestrator, including application health and email delivery metrics.
|
||||
"""
|
||||
|
||||
from prometheus_client import Counter, Gauge
|
||||
|
||||
APP_UP = Gauge(
|
||||
'app_up',
|
||||
'Indicates if the application is running (1) or shutting down (0)',
|
||||
['pod_id'],
|
||||
)
|
||||
|
||||
CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
|
||||
|
||||
|
||||
EMAIL_SENT_COUNT = Counter(
|
||||
'email_sent_count',
|
||||
'Total number of emails sent by the orchestrator',
|
||||
[*CORE_LABELS, 'email_group'],
|
||||
)
|
||||
0
orchestrator/utils/__init__.py
Normal file
0
orchestrator/utils/__init__.py
Normal file
97
orchestrator/utils/connectors_config.py
Normal file
97
orchestrator/utils/connectors_config.py
Normal file
@@ -0,0 +1,97 @@
|
||||
from os import getenv
|
||||
|
||||
|
||||
def build_redis_config():
|
||||
"""
|
||||
Build Redis configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: Redis configuration with host, port, username, and password.
|
||||
"""
|
||||
return {
|
||||
'host': getenv('REDIS_HOST', 'localhost'),
|
||||
'port': int(getenv('REDIS_PORT', '6379')),
|
||||
'username': getenv('REDIS_USERNAME', 'default'),
|
||||
'password': getenv('REDIS_PASSWORD', 'bdnZOpcyiL'),
|
||||
}
|
||||
|
||||
|
||||
def build_mongodb_config():
|
||||
"""
|
||||
Build MongoDB configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: MongoDB configuration with connection string, database name, and TTL index seconds.
|
||||
"""
|
||||
username = getenv('MONGODB_USERNAME', 'root')
|
||||
password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
|
||||
uri = getenv('MONGODB_URL', 'localhost:27018')
|
||||
|
||||
connection_string = f'mongodb://{username}:{password}@{uri}'
|
||||
return {
|
||||
'connection_string': connection_string,
|
||||
'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
|
||||
'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
|
||||
}
|
||||
|
||||
|
||||
def build_couchbase_config():
|
||||
"""
|
||||
Build Couchbase configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: Couchbase configuration with connection string, username, and password.
|
||||
"""
|
||||
return {
|
||||
'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'),
|
||||
'username': getenv('COUCHBASE_USERNAME', 'sientia'),
|
||||
'password': getenv('COUCHBASE_PASSWORD', 'sientia'),
|
||||
}
|
||||
|
||||
|
||||
def build_temporal_config():
|
||||
"""
|
||||
Build Temporal configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: Temporal configuration with host and namespace settings.
|
||||
"""
|
||||
return {
|
||||
'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'),
|
||||
'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'),
|
||||
'temporal_scouter_namespace': getenv('TEMPORAL_SCOUTER_NAMESPACE', 'scouter'),
|
||||
'temporal_laborious_namespace': getenv('TEMPORAL_LABORIOUS_NAMESPACE', 'laborious'),
|
||||
}
|
||||
|
||||
|
||||
def build_postgres_config():
|
||||
"""
|
||||
Build PostgreSQL configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: PostgreSQL configuration with connection details and connection pool settings.
|
||||
"""
|
||||
return {
|
||||
'host': getenv('POSTGRES_HOST', 'localhost'),
|
||||
'port': int(getenv('POSTGRES_PORT', '5432')),
|
||||
'user': getenv('POSTGRES_USER', 'sientia'),
|
||||
'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')),
|
||||
}
|
||||
|
||||
|
||||
def build_email_config():
|
||||
"""
|
||||
Build email configuration from environment variables.
|
||||
|
||||
Returns:
|
||||
dict: Email configuration with SMTP server settings and sender credentials.
|
||||
"""
|
||||
return {
|
||||
'sender_email': getenv('EMAIL_SENDER', 'sientia-alerts@aignosi.com'),
|
||||
'sender_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'),
|
||||
'smtp_server': getenv('EMAIL_SMTP_SERVER', None),
|
||||
'smtp_port': int(getenv('EMAIL_SMTP_PORT', '587')),
|
||||
}
|
||||
31
orchestrator/utils/converters.py
Normal file
31
orchestrator/utils/converters.py
Normal file
@@ -0,0 +1,31 @@
|
||||
def parse_frequency(frequency: str) -> int:
|
||||
"""
|
||||
Parse frequency string into seconds for Temporal schedule intervals.
|
||||
|
||||
This function converts human-readable frequency strings into seconds
|
||||
for use in Temporal schedule configurations. Supports seconds, minutes,
|
||||
hours, and days notation.
|
||||
|
||||
Args:
|
||||
frequency (str): Frequency string with suffix:
|
||||
- 's' for seconds (e.g., '30s')
|
||||
- 'm' for minutes (e.g., '5m')
|
||||
- 'h' for hours (e.g., '2h')
|
||||
- 'd' for days (e.g., '1d')
|
||||
|
||||
Returns:
|
||||
int: Frequency converted to seconds
|
||||
|
||||
Raises:
|
||||
ValueError: If frequency format is invalid
|
||||
"""
|
||||
if frequency.endswith('s'):
|
||||
return int(frequency[:-1])
|
||||
elif frequency.endswith('m'):
|
||||
return int(frequency[:-1]) * 60
|
||||
elif frequency.endswith('h'):
|
||||
return int(frequency[:-1]) * 60 * 60
|
||||
elif frequency.endswith('d'):
|
||||
return int(frequency[:-1]) * 60 * 60 * 24
|
||||
else:
|
||||
raise ValueError('Invalid frequency')
|
||||
132
orchestrator/utils/email_builder.py
Normal file
132
orchestrator/utils/email_builder.py
Normal file
@@ -0,0 +1,132 @@
|
||||
from typing import Any
|
||||
|
||||
from jinja2 import Template
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
|
||||
class EmailBuilder:
|
||||
"""
|
||||
HTML email template builder for notification emails.
|
||||
|
||||
This class handles the generation of HTML email content from notification
|
||||
data using Jinja2 templates. It supports different email types (alerts,
|
||||
reports) and notification levels (ERROR, WARNING, INFO) with customizable
|
||||
templates and parameter replacement.
|
||||
|
||||
Args:
|
||||
logger (Logger): Application logger instance for error reporting
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Logger):
|
||||
self.logger = logger
|
||||
|
||||
self.report_template_file = './orchestrator/utils/templates/email_template.html'
|
||||
self.general_template_file = './orchestrator/utils/templates/general_template.html'
|
||||
|
||||
with open(self.report_template_file) as file:
|
||||
self.report_template = file.read()
|
||||
with open(self.general_template_file) as file:
|
||||
self.general_template = file.read()
|
||||
|
||||
def replace_parameters(self, template: str, parameters: dict) -> str:
|
||||
"""
|
||||
Replace parameters in a Jinja2 template with provided values.
|
||||
|
||||
Renders a Jinja2 template string with the provided parameter dictionary,
|
||||
replacing all template variables with their corresponding values.
|
||||
|
||||
Args:
|
||||
template (str): The Jinja2 template string
|
||||
parameters (dict): Dictionary of parameters to replace in the template
|
||||
|
||||
Returns:
|
||||
str: The rendered template with parameters replaced
|
||||
"""
|
||||
# Create a Jinja2 template from the provided string
|
||||
template_obj = Template(template)
|
||||
|
||||
return template_obj.render(parameters)
|
||||
|
||||
def parameters(self, general_events: dict, mail_type: str) -> dict:
|
||||
"""
|
||||
Build parameters dictionary for email templates based on general events and mail type.
|
||||
|
||||
Processes notification events organized by level and model, rendering HTML
|
||||
sections for each notification level using the general template.
|
||||
|
||||
Args:
|
||||
general_events (dict): Dictionary containing events categorized by level (ERROR, WARNING, INFO).
|
||||
Each level contains a 'models' key with model-specific event data
|
||||
mail_type (str): The type of email being sent (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
dict: Dictionary with mail_type and rendered event sections for each notification level
|
||||
"""
|
||||
error_events = general_events.get('ERROR', {})
|
||||
warning_events = general_events.get('WARNING', {})
|
||||
info_events = general_events.get('INFO', {})
|
||||
|
||||
error_models = error_events.get('models', [])
|
||||
warning_models = warning_events.get('models', [])
|
||||
info_models = info_events.get('models', [])
|
||||
|
||||
return {
|
||||
'mail_type': mail_type,
|
||||
'error_events': self.replace_parameters(self.general_template, error_events)
|
||||
if error_models
|
||||
else '',
|
||||
'warning_events': self.replace_parameters(self.general_template, warning_events)
|
||||
if warning_models
|
||||
else '',
|
||||
'info_events': self.replace_parameters(self.general_template, info_events)
|
||||
if info_models
|
||||
else '',
|
||||
}
|
||||
|
||||
def build_email(self, report_data: list[dict[str, Any]], mail_type: str) -> str:
|
||||
"""
|
||||
Build the email HTML by organizing report data by notification level and model.
|
||||
|
||||
Organizes notification data by level and model, then renders the complete
|
||||
HTML email using the report template with all event sections.
|
||||
|
||||
Args:
|
||||
report_data (list[dict[str, Any]]): List of notification reports, each containing:
|
||||
- level (str): Notification level (ERROR, WARNING, INFO)
|
||||
- model_name (str): Name of the model
|
||||
- Additional notification details
|
||||
mail_type (str): The type of email being built (Alerts/Reports)
|
||||
|
||||
Returns:
|
||||
str: Complete HTML email content ready for sending
|
||||
"""
|
||||
general_events: dict[str, dict[str, Any]] = {}
|
||||
|
||||
for report in report_data:
|
||||
level = report['level']
|
||||
model_name = report['model_name']
|
||||
|
||||
if level not in general_events:
|
||||
general_events[level] = {
|
||||
'section_name': f'{level.capitalize()}s detected:',
|
||||
'models': {},
|
||||
}
|
||||
|
||||
# Type assertion to help the type checker understand the structure
|
||||
level_data = general_events[level]
|
||||
models_dict = level_data['models']
|
||||
|
||||
if model_name not in models_dict:
|
||||
models_dict[model_name] = {
|
||||
'model_name': model_name,
|
||||
'events': [],
|
||||
}
|
||||
|
||||
models_dict[model_name]['events'].append(report)
|
||||
|
||||
for _type, content in general_events.items():
|
||||
content['models'] = list(content['models'].values())
|
||||
|
||||
return self.replace_parameters(
|
||||
self.report_template, self.parameters(general_events, mail_type)
|
||||
)
|
||||
505
orchestrator/utils/orchestrator_functions.py
Normal file
505
orchestrator/utils/orchestrator_functions.py
Normal file
@@ -0,0 +1,505 @@
|
||||
from typing import Any
|
||||
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
|
||||
def common_config(config: dict[str, Any]):
|
||||
"""
|
||||
Extract common configuration parameters from a pipeline configuration.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- workflow_type (str): Type of workflow (e.g., 'scouter', 'predictions_batch', 'drift')
|
||||
- schedule_name (str): Unique name identifier for this schedule
|
||||
- model_id (str): MongoDB ID of the associated model
|
||||
- model (dict): Model configuration containing:
|
||||
- name (str): Human-readable name of the model
|
||||
- model_config (dict, optional): Additional model-specific configuration
|
||||
- frequency (str, optional): Execution frequency (default: '1m')
|
||||
Format: '{number}{unit}' where unit is 's', 'm', 'h', or 'd'
|
||||
- offset (str, optional): Schedule offset/delay (default: '0m')
|
||||
- max_retry_policy (int, optional): Maximum retry attempts on failure (default: 1)
|
||||
- execution_timeout_seconds (int, optional): Workflow execution timeout (default: 300)
|
||||
- task_timeout_seconds (int, optional): Individual task timeout (default: 300)
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Common configuration dictionary with standardized parameters
|
||||
for Temporal workflow execution
|
||||
"""
|
||||
model = config['model']
|
||||
return {
|
||||
'workflow_type': config['workflow_type'],
|
||||
'schedule_name': config['schedule_name'],
|
||||
'frequency': config.get('frequency', '1m'),
|
||||
'offset': config.get('offset', '0m'),
|
||||
'max_retry_policy': config.get('max_retry_policy', 1),
|
||||
'model_id': config['model_id'],
|
||||
'model_name': model['name'],
|
||||
'model_config': model.get('model_config', {}),
|
||||
'execution_timeout_seconds': config.get('execution_timeout_seconds', 300),
|
||||
'task_timeout_seconds': config.get('task_timeout_seconds', 300),
|
||||
'on_conflict': config.get('on_conflict', 'error'),
|
||||
'runtime': config.get('runtime', 'legacy'),
|
||||
}
|
||||
|
||||
|
||||
def drift(config: dict[str, Any]):
|
||||
"""
|
||||
Build drift configuration from pipeline config.
|
||||
|
||||
Creates a drift detection workflow configuration with database table mappings
|
||||
and drift metric specifications for monitoring data distribution changes.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- interval_minutes (int, optional): Time window for data comparison in minutes (default: 60)
|
||||
- drift_metrics (list[str], optional): Statistical metrics to compute (default:
|
||||
['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein'])
|
||||
- Additional fields from common_config
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Drift detection configuration with source/target tables,
|
||||
time interval, and metrics specifications. Results stored in 'drift_metrics' table
|
||||
"""
|
||||
return {
|
||||
**common_config(config),
|
||||
'schema': 'sientia_data',
|
||||
'source_table_name': 'laborious_data',
|
||||
'target_table_name': 'drift_metrics',
|
||||
'interval': config.get('interval_minutes', 60),
|
||||
'drift_metrics': config.get(
|
||||
'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def simple_metrics(config: dict[str, Any]):
|
||||
"""
|
||||
Build simple metrics configuration from pipeline config.
|
||||
|
||||
Creates a simple metrics computation workflow configuration for calculating
|
||||
model performance metrics like RMSE, MSE, MAE, and R².
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- interval_minutes (int, optional): Computation interval in minutes (default: 60)
|
||||
- metrics (list[str], optional): List of metrics to compute
|
||||
(default: ['rmse', 'mse', 'mae', 'r2'])
|
||||
- Additional fields from common_config
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Simple metrics configuration with source tables (predictions and
|
||||
actual data), target table, time interval, and metrics list. Results stored in
|
||||
'simple_metrics' table
|
||||
"""
|
||||
return {
|
||||
**common_config(config),
|
||||
'schema': 'sientia_data',
|
||||
'predictions_table_name': 'predictions',
|
||||
'data_table_name': 'laborious_data',
|
||||
'target_table_name': 'simple_metrics',
|
||||
'interval_minutes': config.get('interval_minutes', 60),
|
||||
'metrics': config.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
|
||||
}
|
||||
|
||||
|
||||
def minimal_retrain(config: dict[str, Any]):
|
||||
"""
|
||||
Build minimal retrain configuration from pipeline config.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- query (str): SQL query to retrieve training data. Should return features
|
||||
and target variable in expected format
|
||||
- datetime_columns (list[str], optional): Column names to parse as datetime
|
||||
for proper temporal handling (default: [])
|
||||
- Additional fields from common_config
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Minimal retrain configuration with SQL query, database settings,
|
||||
and datetime column specifications. Retraining logs are stored in 'log_retrain' table
|
||||
"""
|
||||
return {
|
||||
**common_config(config),
|
||||
'query': config['query'],
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_retrain',
|
||||
'datetime_columns': config.get('datetime_columns', []),
|
||||
}
|
||||
|
||||
|
||||
def base_scouter(config: dict[str, Any]):
|
||||
"""
|
||||
Build base scouter configuration shared by all scouter workflow types.
|
||||
|
||||
Creates the foundational configuration for OPC data collection workflows,
|
||||
including filter policies, database settings, and data retention parameters.
|
||||
This configuration is extended by specific scouter implementations (OPC UA, PI Web API).
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- filters (list[dict], optional): List of filter configurations with:
|
||||
- filter_name (str): Name of the filter
|
||||
- policy (str): Filter policy to apply
|
||||
- tag_retention_minutes (int, optional): Tag retention time in minutes (default: 60)
|
||||
- debug_data_package (bool, optional): Enable debug data package logging (default: False)
|
||||
- fill_missing_tags (bool, optional): Fill missing tags with interpolation (default: False)
|
||||
- Additional fields from common_config
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Base scouter configuration with filters, database settings,
|
||||
and retention policies
|
||||
"""
|
||||
|
||||
filters = {}
|
||||
for f in config.get('filters', []):
|
||||
filters[f['filter_name']] = {'policy': f['policy']}
|
||||
|
||||
return {
|
||||
**common_config(config),
|
||||
'trigger_laborious': False,
|
||||
'filters': filters,
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': config.get('tag_retention_minutes', 60) * 60,
|
||||
'debug_data_package': config.get('debug_data_package', False),
|
||||
'fill_missing_tags': config.get('fill_missing_tags', False),
|
||||
}
|
||||
|
||||
|
||||
def scouter(config: dict[str, Any]):
|
||||
"""
|
||||
Build scouter configuration from pipeline config.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- schedule_name (str): Name of the schedule (used for topic generation)
|
||||
- read_tags (list[dict]): List of tag configurations with:
|
||||
- tag_name (str): Name of the tag to read
|
||||
- aggr_func (str, optional): Aggregation function for data collection (default: 'lts')
|
||||
Common values: 'lts' (last), 'avg' (average), 'min', 'max', 'sum'
|
||||
- data_range (list[int], optional): Valid data range [min, max] (default: [-100, 100])
|
||||
- Additional fields from base_scouter
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: OPC UA scouter configuration with Kafka topic, tag mappings,
|
||||
filters, and retention settings. Topic name follows pattern: 'raw_{schedule_name}'
|
||||
"""
|
||||
|
||||
tags = {}
|
||||
for tag in config['read_tags']:
|
||||
tags[tag['tag_name']] = {
|
||||
'aggr_func': tag.get('aggr_func', 'lts'),
|
||||
'data_range': tag.get('data_range', [-100, 100]),
|
||||
}
|
||||
|
||||
return {
|
||||
**base_scouter(config),
|
||||
'topic': f'raw_{config["schedule_name"]}',
|
||||
'model_tags': tags,
|
||||
}
|
||||
|
||||
|
||||
def pi_web_api_scouter(config: dict[str, Any]):
|
||||
"""
|
||||
Build PI Web API scouter configuration from pipeline config.
|
||||
|
||||
Creates a data collection workflow configuration for OSIsoft PI servers using
|
||||
the PI Web API REST interface. Configures tag mappings with WebIDs, aggregation
|
||||
functions, and API query parameters including timeout management.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- read_tags (list[dict]): List of tag configurations with:
|
||||
- tag_name (str): Name of the tag
|
||||
- webid (str): PI Web API WebID for the tag
|
||||
- aggr_func (str, optional): Aggregation function (default: 'lts')
|
||||
- data_range (list[int], optional): Valid data range (default: [-100, 100])
|
||||
- pi_web_api_config (dict): PI Web API connection settings with:
|
||||
- endpoint (str): PI Web API endpoint URL
|
||||
- period (str, optional): Time period for data retrieval (default: '*-1d')
|
||||
- max_count (int, optional): Maximum number of values to retrieve (default: 1)
|
||||
- api_timeout (int, optional): API request timeout in seconds
|
||||
- Additional fields from base_scouter
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: PI Web API scouter configuration with tag mappings and query settings.
|
||||
API timeout is automatically adjusted to not exceed workflow frequency.
|
||||
"""
|
||||
tags = {}
|
||||
for tag, tag_config in config['read_tags'].items():
|
||||
tags[tag] = {
|
||||
'webid': tag_config['webid'],
|
||||
'aggr_func': tag_config.get('aggr_func', 'lts'),
|
||||
'data_range': tag_config.get('data_range', [-100, 100]),
|
||||
}
|
||||
|
||||
base_config = base_scouter(config)
|
||||
|
||||
pi_web_api_config = config['pi_web_api_config']
|
||||
|
||||
config_timeout = pi_web_api_config.get('api_timeout', None)
|
||||
frequency = parse_frequency(base_config['frequency'])
|
||||
|
||||
if config_timeout is None or config_timeout > frequency:
|
||||
config_timeout = frequency
|
||||
|
||||
return {
|
||||
**base_config,
|
||||
'model_tags': tags,
|
||||
'pi_web_api_query': {
|
||||
'endpoint': pi_web_api_config['endpoint'],
|
||||
'period': pi_web_api_config.get('period', '*-1d'),
|
||||
'max_count': pi_web_api_config.get('max_count', 1),
|
||||
'api_timeout': config_timeout,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[str, Any]]):
|
||||
"""
|
||||
Merge filter configurations with base filter configuration.
|
||||
|
||||
Extends the base filter configuration dictionary by adding or overwriting
|
||||
filters from the provided configuration list. Used in predictions_batch
|
||||
workflows to combine default filters with user-defined custom filters.
|
||||
|
||||
Args:
|
||||
base_filter_config (dict[str, Any]): Base filter configuration dictionary to extend.
|
||||
Each filter entry contains 'policy' and optionally 'config' keys.
|
||||
config (list[dict[str, Any]]): List of filter configurations to merge, each containing:
|
||||
- filter_name (str): Name of the filter to add or update
|
||||
- policy (str): Filter policy (e.g., 'STOP', 'CONTINUE', 'REPEAT')
|
||||
- config (dict, optional): Additional filter-specific configuration
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Extended filter configuration dictionary with merged filters.
|
||||
Filters from config list overwrite or add to base_filter_config entries.
|
||||
"""
|
||||
for fil in config:
|
||||
base_filter_config[fil['filter_name']] = {
|
||||
'policy': fil['policy'],
|
||||
'config': fil.get('config', {}),
|
||||
}
|
||||
|
||||
return base_filter_config
|
||||
|
||||
|
||||
def process_path_priority(path_priority: list[str]):
|
||||
"""
|
||||
Process and normalize path priority list for filter policy execution order.
|
||||
|
||||
Validates and normalizes the priority list used to determine the order in which
|
||||
filter policies are evaluated in prediction workflows. Invalid priorities are
|
||||
removed, missing required priorities are appended, and the list is truncated to
|
||||
exactly 3 elements.
|
||||
|
||||
Valid priorities define workflow behavior when filters are triggered:
|
||||
- STOP: Halt workflow execution immediately
|
||||
- CONTINUE: Proceed to next step despite filter trigger
|
||||
- REPEAT: Retry the current step
|
||||
|
||||
Args:
|
||||
path_priority (list[str]): User-provided list of path priorities. May contain
|
||||
invalid values or be incomplete.
|
||||
|
||||
Returns:
|
||||
list[str]: Normalized path priority list with exactly 3 elements in user-specified
|
||||
or default order. Default order when priorities are missing: ["STOP", "CONTINUE", "REPEAT"]
|
||||
"""
|
||||
for priority in path_priority[:]:
|
||||
if priority not in ['STOP', 'CONTINUE', 'REPEAT']:
|
||||
path_priority.remove(priority)
|
||||
|
||||
for priority in ['STOP', 'CONTINUE', 'REPEAT']:
|
||||
if priority not in path_priority:
|
||||
path_priority.append(priority)
|
||||
|
||||
return path_priority[0:3]
|
||||
|
||||
|
||||
def predictions_batch(config: dict[str, Any]):
|
||||
"""
|
||||
Build predictions batch configuration from pipeline config.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Pipeline configuration containing:
|
||||
- query (str): SQL query to retrieve input data for predictions
|
||||
- write_tags (list[dict]): List of OPC tag configurations for write-back with:
|
||||
- server_id (str): ID of the target OPC server
|
||||
- type (str): Tag type - 'prediction' (model output) or 'confidence' (prediction confidence)
|
||||
- addr (str): OPC tag address/path
|
||||
- data_type (str, optional): OPC data type (default: 'float')
|
||||
- datetime_columns (list[str], optional): Column names to parse as datetime (default: [])
|
||||
- path_priority (list[str], optional): Filter policy execution order (default: ["STOP", "CONTINUE", "REPEAT"])
|
||||
- input_filters (list[dict], optional): Input data validation filters
|
||||
- mlflow_transform_filters (list[dict], optional): Transform stage filters
|
||||
- mlflow_predict_filters (list[dict], optional): Prediction stage filters
|
||||
- model_retention_minutes (int, optional): Data retention time in minutes (default: 60)
|
||||
- save_transform (bool, optional): Save transformed data to database (default: True)
|
||||
- predictions_storage_policy (str, optional): Prediction storage policy (default: 'lts:1')
|
||||
- pi_web_api_output_config (dict, optional): PI Web API output configuration for write-back (default: {})
|
||||
- Additional fields from common_config
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Complete predictions batch configuration with OPC output mappings,
|
||||
PI Web API output configuration, multi-stage filters, SQL query, and retention policies
|
||||
"""
|
||||
tags: dict[str, Any] = {}
|
||||
for tag in config.get('write_tags', []):
|
||||
if tag['server_id'] not in tags:
|
||||
tags[tag['server_id']] = {}
|
||||
|
||||
tag_type = tag['type']
|
||||
|
||||
if tag_type == 'prediction' or tag_type == 'confidence':
|
||||
tag_type_str = f'{tag_type}_tags'
|
||||
|
||||
if tag_type_str not in tags[tag['server_id']]:
|
||||
tags[tag['server_id']][tag_type_str] = {}
|
||||
|
||||
tags[tag['server_id']][tag_type_str][tag['addr']] = {
|
||||
'data_type': tag.get('data_type', 'float'),
|
||||
}
|
||||
|
||||
path_priority = process_path_priority(
|
||||
config.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT'])
|
||||
)
|
||||
|
||||
return {
|
||||
**common_config(config),
|
||||
'query': config['query'],
|
||||
'datetime_columns': config.get('datetime_columns', []),
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'predictions',
|
||||
'save_transform': config.get('save_transform', True),
|
||||
'transform_table_name': 'transformed_data',
|
||||
'retention_time': config.get('model_retention_minutes', 60) * 60,
|
||||
'opc_output_config': tags,
|
||||
'pi_web_api_output_config': config.get('pi_web_api_output_config', {}),
|
||||
'input_filters': overlap_filter_config(
|
||||
{'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config.get('input_filters', [])
|
||||
),
|
||||
'mlflow_transform_filters': overlap_filter_config(
|
||||
{
|
||||
'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
|
||||
'API_ERROR': {'policy': 'STOP', 'config': {}},
|
||||
},
|
||||
config.get('mlflow_transform_filters', []),
|
||||
),
|
||||
'mlflow_predict_filters': overlap_filter_config(
|
||||
{'API_ERROR': {'policy': 'STOP', 'config': {}}},
|
||||
config.get('mlflow_predict_filters', []),
|
||||
),
|
||||
'path_priority': path_priority,
|
||||
'predictions_storage_policy': config.get('predictions_storage_policy', 'lts:1'),
|
||||
}
|
||||
|
||||
|
||||
def gather_read_tags(pipelines: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
"""
|
||||
Gather all read tags from scouter pipeline configurations.
|
||||
|
||||
Collects all read tags from scouter-type pipelines and organizes them by
|
||||
server_id and tag_address, tracking which Kafka topics each tag is associated with.
|
||||
This function is used during slot configuration to aggregate tags across multiple
|
||||
scouter pipelines for efficient OPC server slot allocation.
|
||||
|
||||
Args:
|
||||
pipelines (list[dict[str, Any]]): List of pipeline configurations to process.
|
||||
Only pipelines with workflow_type 'scouter' are processed. Each scouter
|
||||
pipeline should contain a 'read_tags' list with tag configurations.
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Dictionary of read tags keyed by "server_id:tag_address",
|
||||
where each entry contains:
|
||||
- All original tag configuration fields
|
||||
- topics (list[str]): List of Kafka topic names associated with this tag
|
||||
(format: 'raw_{schedule_name}')
|
||||
"""
|
||||
|
||||
tags = {}
|
||||
|
||||
# Get all read tags from pipelines
|
||||
for pipeline in pipelines:
|
||||
if pipeline['workflow_type'] != 'scouter':
|
||||
continue
|
||||
|
||||
for tag in pipeline.get('read_tags', []):
|
||||
tag_string = f'{tag["server_id"]}:{tag["tag_address"]}'
|
||||
if tag_string not in tags:
|
||||
tags[tag_string] = {**tag, 'topics': []}
|
||||
|
||||
tags[tag_string]['topics'].append(f'raw_{pipeline["schedule_name"]}')
|
||||
|
||||
return tags
|
||||
|
||||
|
||||
def build_tag_config(
|
||||
tags: list[dict[str, Any]], opc_servers: dict[str, Any]
|
||||
) -> tuple[dict[str, Any], list]:
|
||||
"""
|
||||
Build tag configuration for a specific slot and OPC server.
|
||||
|
||||
Organizes tags by OPC server name and calculates the minimum subscription period
|
||||
based on tag frequencies. Validates that all server IDs exist in the OPC
|
||||
servers configuration. The subscription period is set to half of the minimum
|
||||
tag frequency to ensure efficient data collection.
|
||||
|
||||
Args:
|
||||
tags (list[dict[str, Any]]): List of tag configurations containing:
|
||||
- server_id (str): ID of the OPC server
|
||||
- tag_address (str): Address/path of the OPC tag
|
||||
- frequency (int): Tag read frequency in milliseconds
|
||||
- Additional tag-specific configuration fields
|
||||
opc_servers (dict[str, Any]): Dictionary of OPC server configurations keyed by server_id.
|
||||
Each server configuration should contain:
|
||||
- server_name (str): Human-readable server name
|
||||
- url (str): OPC server URL
|
||||
- uri (str): OPC server URI
|
||||
- cert_path (str, optional): Certificate file path
|
||||
- private_key_path (str, optional): Private key file path
|
||||
- server_cert_path (str, optional): Server certificate file path
|
||||
|
||||
Returns:
|
||||
tuple[dict[str, Any], list]: A tuple containing:
|
||||
- Slot configuration dictionary organized by server_name, where each server
|
||||
contains connection details, tags dictionary, and subscription_period_ms
|
||||
- List of server IDs (str) that were not found in opc_servers configuration
|
||||
"""
|
||||
|
||||
slot_config = {}
|
||||
notifications = []
|
||||
|
||||
for tag in tags:
|
||||
server_id = tag['server_id']
|
||||
|
||||
if server_id not in opc_servers:
|
||||
notifications.append(server_id)
|
||||
continue
|
||||
|
||||
server_name = opc_servers[server_id]['server_name']
|
||||
if server_name not in slot_config:
|
||||
slot_config[server_name] = {
|
||||
'server_id': server_id,
|
||||
'name': server_name,
|
||||
'url': opc_servers[server_id]['url'],
|
||||
'server_uri': opc_servers[server_id]['uri'],
|
||||
'cert_path': opc_servers[server_id].get('cert_path', None),
|
||||
'private_key_path': opc_servers[server_id].get('private_key_path', None),
|
||||
'server_cert_path': opc_servers[server_id].get('server_cert_path', None),
|
||||
'tags': {},
|
||||
}
|
||||
|
||||
slot_config[server_name]['tags'][tag['tag_address']] = {
|
||||
**tag,
|
||||
}
|
||||
|
||||
for server_name in slot_config:
|
||||
frequencies = [int(x['frequency']) for x in slot_config[server_name]['tags'].values()]
|
||||
|
||||
min_frequency = min(frequencies) if frequencies else 1000
|
||||
|
||||
slot_config[server_name]['subscription_period_ms'] = min_frequency / 2
|
||||
|
||||
return slot_config, notifications
|
||||
24
orchestrator/utils/templates/email_template.html
Normal file
24
orchestrator/utils/templates/email_template.html
Normal file
@@ -0,0 +1,24 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="pt-BR">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>SIENTIA™ Report</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; margin: 0; padding: 20px; }
|
||||
h1, h2 { color: #333; }
|
||||
table { width: 100%; border-collapse: collapse; margin-top: 20px; }
|
||||
th, td { border: 1px solid #ddd; padding: 8px; text-align: left; }
|
||||
th { background-color: #f4f4f4; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>SIENTIA™ {{ mail_type }}</h1>
|
||||
|
||||
{{ error_events }}
|
||||
{{ warning_events }}
|
||||
{{ info_events }}
|
||||
|
||||
{{ special_events }}
|
||||
</body>
|
||||
</html>
|
||||
26
orchestrator/utils/templates/general_template.html
Normal file
26
orchestrator/utils/templates/general_template.html
Normal file
@@ -0,0 +1,26 @@
|
||||
<h3>{{ section_name }}</h3>
|
||||
{% for model in models %}
|
||||
<h4>Model: <span>{{ model.model_name }}</span></h4>
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Notification ID</th>
|
||||
<th>Schedule</th>
|
||||
<th>Block</th>
|
||||
<th>Timestamp</th>
|
||||
<th>Message</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for event in model.events %}
|
||||
<tr>
|
||||
<td>{{ event.notification_id }}</td>
|
||||
<td>{{ event.trigger }}</td>
|
||||
<td>{{ event.block }}</td>
|
||||
<td>{{ event.timestamp }}</td>
|
||||
<td>{{ event.message }}</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
{% endfor %}
|
||||
0
orchestrator/worker/__init__.py
Normal file
0
orchestrator/worker/__init__.py
Normal file
228
orchestrator/worker/worker.py
Normal file
228
orchestrator/worker/worker.py
Normal file
@@ -0,0 +1,228 @@
|
||||
from temporalio import client, workflow
|
||||
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
|
||||
from prometheus_client import start_http_server
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import get_logger
|
||||
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
||||
|
||||
from orchestrator import metrics
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.utils.connectors_config import (
|
||||
build_email_config,
|
||||
build_mongodb_config,
|
||||
build_postgres_config,
|
||||
build_redis_config,
|
||||
build_temporal_config,
|
||||
)
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
from orchestrator.workflows.reports import Reports
|
||||
from orchestrator.workflows.subworkflows.load_notification_package import (
|
||||
LoadNotificationPackage,
|
||||
)
|
||||
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
|
||||
|
||||
POD_ID = os.getenv('POD_ID')
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
Main function to initialize and run the Temporal worker.
|
||||
|
||||
Sets up MongoDB connection, notification handler, Temporal client, and starts
|
||||
multiple workers for different task queues (orchestrator, alerts, reports).
|
||||
Handles graceful shutdown and error handling. Initializes Prometheus metrics
|
||||
server and SDK metrics for monitoring.
|
||||
"""
|
||||
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
|
||||
namespace = os.getenv('TEMPORAL_NAMESPACE', 'default')
|
||||
logger = get_logger(__name__)
|
||||
|
||||
metadata = {
|
||||
'pod_id': POD_ID,
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
'workflow_name': '-',
|
||||
'schedule_name': '-',
|
||||
}
|
||||
|
||||
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata=metadata)
|
||||
|
||||
logger.custom_info('Starting prometheus client...', metadata=metadata)
|
||||
start_prometheus_server()
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata=metadata)
|
||||
|
||||
mongo_config = build_mongodb_config()
|
||||
notification_handler = NotificationHandler(
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database=mongo_config['database_name'],
|
||||
logger=logger,
|
||||
project_name=os.getenv('PROJECT_NAME', 'orchestrator'),
|
||||
)
|
||||
|
||||
logger.custom_info(
|
||||
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata=metadata
|
||||
)
|
||||
|
||||
new_runtime = Runtime(
|
||||
telemetry=TelemetryConfig(
|
||||
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
|
||||
)
|
||||
)
|
||||
|
||||
logger.custom_info(f'Starting Temporal Client at {host}:{namespace}', metadata=metadata)
|
||||
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=host,
|
||||
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
|
||||
runtime=new_runtime,
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Activities...', metadata=metadata)
|
||||
|
||||
activities = Activities(
|
||||
temporal_config=build_temporal_config(),
|
||||
redis_config=build_redis_config(),
|
||||
mongodb_config=build_mongodb_config(),
|
||||
email_config=build_email_config(),
|
||||
postgres_config=build_postgres_config(),
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
await activities.connect_to_temporal()
|
||||
|
||||
logger.custom_info('Starting Workers...', metadata=metadata)
|
||||
|
||||
workers = [
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
main_workflow=Orchestrator,
|
||||
other_workflows=[],
|
||||
activities=[
|
||||
# Redis
|
||||
activities.load_active_ingestors,
|
||||
activities.load_opc_slots,
|
||||
activities.update_slots,
|
||||
activities.delete_slots,
|
||||
# Couchbase
|
||||
# activities.load_query_from_couchbase,
|
||||
# MongoDB
|
||||
activities.aggregate_documents_in_mongodb,
|
||||
activities.find_documents_in_mongodb,
|
||||
activities.update_pipelines_timestamps,
|
||||
activities.create_pipelines_timestamps,
|
||||
activities.delete_pipelines_timestamps,
|
||||
activities.create_collection_with_ttl_index,
|
||||
# Temporal
|
||||
activities.create_schedules,
|
||||
activities.update_schedules,
|
||||
activities.delete_schedules,
|
||||
activities.normalize_schedules,
|
||||
# Formatters
|
||||
activities.process_schedules,
|
||||
activities.process_slots,
|
||||
activities.create_schedule_config,
|
||||
activities.create_slot_config,
|
||||
activities.report_schedule_orchestration,
|
||||
activities.report_slot_orchestration,
|
||||
activities.format_schedule_config,
|
||||
],
|
||||
logger=logger,
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
main_workflow=Alerts,
|
||||
other_workflows=[LoadNotificationPackage, ProcessNotifications],
|
||||
activities=[
|
||||
# Load notifications
|
||||
activities.get_last_data_timestamp,
|
||||
activities.find_documents_in_mongodb,
|
||||
activities.load_latest_data,
|
||||
activities.put_last_data_timestamp,
|
||||
# Format and filter notifications
|
||||
activities.filter_notification_alerts,
|
||||
# Send email and export data to postgres
|
||||
activities.build_email_html,
|
||||
activities.send_email,
|
||||
activities.format_log_report,
|
||||
activities.export_data_to_postgres,
|
||||
# Store notification cache
|
||||
activities.store_notification_cache,
|
||||
],
|
||||
logger=logger,
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
main_workflow=Reports,
|
||||
other_workflows=[LoadNotificationPackage, ProcessNotifications],
|
||||
activities=[
|
||||
# Load notifications
|
||||
activities.get_last_data_timestamp,
|
||||
activities.find_documents_in_mongodb,
|
||||
activities.load_latest_data,
|
||||
activities.put_last_data_timestamp,
|
||||
# Format and filter notifications
|
||||
activities.filter_notification_reports,
|
||||
# Send email and export data to postgres
|
||||
activities.build_email_html,
|
||||
activities.send_email,
|
||||
activities.format_log_report,
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
),
|
||||
]
|
||||
|
||||
handlers = []
|
||||
for w in workers:
|
||||
handlers.append(w.run())
|
||||
|
||||
logger.custom_info('Workers started successfully', metadata=metadata)
|
||||
|
||||
exit_code = 0
|
||||
try:
|
||||
await asyncio.gather(*handlers)
|
||||
except BaseException as e:
|
||||
logger.error(f'An unhandled exception occurred: {e}', exc_info=True)
|
||||
exit_code = 1
|
||||
finally:
|
||||
if notification_handler:
|
||||
notification_handler.shutdown()
|
||||
if activities:
|
||||
activities.shutdown()
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
||||
sys.exit(exit_code)
|
||||
|
||||
|
||||
def start_prometheus_server():
|
||||
"""
|
||||
Start the Prometheus metrics server on the configured port.
|
||||
|
||||
Sets up HTTP server for metrics collection and marks the application as UP.
|
||||
Exits the application if the server fails to start. The metrics server
|
||||
exposes application metrics on the port specified by HTTP_METRICS_PORT.
|
||||
|
||||
Raises:
|
||||
SystemExit: If the metrics server fails to start
|
||||
"""
|
||||
try:
|
||||
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
|
||||
start_http_server(port)
|
||||
print(f'Prometheus server started on port {port}.')
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(1)
|
||||
except Exception as e:
|
||||
print(f'Failed to start Prometheus server: {e}')
|
||||
os._exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
||||
0
orchestrator/workflows/__init__.py
Normal file
0
orchestrator/workflows/__init__.py
Normal file
116
orchestrator/workflows/alerts.py
Normal file
116
orchestrator/workflows/alerts.py
Normal file
@@ -0,0 +1,116 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='alerts')
|
||||
class Alerts:
|
||||
"""
|
||||
Alerts workflow for real-time error notification delivery.
|
||||
|
||||
This workflow processes ERROR-level notifications from the notification queue
|
||||
and sends immediate alerts to configured user groups. It implements intelligent
|
||||
filtering with TTL-based duplicate prevention and persistent alert detection
|
||||
for ongoing issues.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the alerts workflow for real-time error notification delivery.
|
||||
|
||||
This workflow loads ERROR-level notifications from MongoDB, applies
|
||||
intelligent filtering with TTL management to prevent alert spam,
|
||||
and sends immediate email alerts to configured receiver groups.
|
||||
|
||||
The workflow implements:
|
||||
- TTL-based duplicate prevention for notifications
|
||||
- Persistent alert detection for ongoing issues
|
||||
- User group-based filtering with ignore lists
|
||||
- Audit logging of alert delivery status
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Workflow input parameters.
|
||||
Required fields:
|
||||
- schedule_name (str): Name of the alert schedule
|
||||
- notification_ttl (int): Seconds before considering notification persistent
|
||||
- sent_ttl (int): Time-to-live for sent notification cache
|
||||
|
||||
Returns:
|
||||
None: Workflow completes without return value
|
||||
|
||||
Raises:
|
||||
Exception: If alert processing or delivery fails
|
||||
"""
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'workflow_name': 'alerts',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
mail_type = 'Alerts'
|
||||
|
||||
input_data['metadata'] = metadata
|
||||
input_data['mail_type'] = mail_type
|
||||
|
||||
input_data['base_data_filter'] = {'level': 'ERROR'}
|
||||
|
||||
# Call subworkflow "load_notification_package" passing the static filters
|
||||
# (level = "ERROR" and timestamp > last timestamp)
|
||||
|
||||
package = await workflow.execute_child_workflow(
|
||||
'subworkflow.load_notification_package', input_data
|
||||
)
|
||||
|
||||
if not package['notification_package'] or not package['sending_configs']:
|
||||
return
|
||||
|
||||
# Filter notification package by groups custom configs, levels and
|
||||
# timestamp cached
|
||||
|
||||
receiver_groups = await workflow.execute_local_activity_method(
|
||||
Activities.filter_notification_alerts,
|
||||
{
|
||||
**metadata,
|
||||
'notification_package': package['notification_package'],
|
||||
'sending_configs': package['sending_configs'],
|
||||
'notification_ttl': input_data['notification_ttl'],
|
||||
},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
if not receiver_groups:
|
||||
return
|
||||
|
||||
# Call subworkflow "process_notifications" passing the notification package
|
||||
log_report = await workflow.execute_child_workflow(
|
||||
'subworkflow.process_notifications',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'mail_type': mail_type,
|
||||
'notification_package': receiver_groups,
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_report',
|
||||
},
|
||||
)
|
||||
|
||||
if not log_report:
|
||||
return
|
||||
|
||||
# Store the notification_id sendings to avoid sending them again
|
||||
await workflow.execute_activity_method(
|
||||
Activities.store_notification_cache,
|
||||
{**metadata, 'log_report': log_report, 'sent_ttl': input_data['sent_ttl']},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
281
orchestrator/workflows/orchestrator.py
Normal file
281
orchestrator/workflows/orchestrator.py
Normal file
@@ -0,0 +1,281 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='orchestrator')
|
||||
class Orchestrator:
|
||||
"""
|
||||
Main orchestrator workflow for pipeline and resource management.
|
||||
|
||||
This workflow coordinates pipeline deployment and OPC server slot
|
||||
management by retrieving configurations from MongoDB and Redis,
|
||||
processing schedules, and deploying them to the Temporal server
|
||||
and Redis infrastructure.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the orchestration workflow for pipeline and slot management.
|
||||
|
||||
This workflow retrieves pipeline configurations and OPC server data,
|
||||
processes schedules and slot configurations, and deploys them to
|
||||
the appropriate services. It handles creation, updates, and deletion
|
||||
of schedules and slots based on current system state.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Workflow input parameters.
|
||||
Required fields:
|
||||
- schedule_name (str): Name of the orchestration schedule
|
||||
- pipelines_query (dict[str, Any]): MongoDB query for pipeline configurations
|
||||
- opc_servers_query (dict[str, Any]): MongoDB query for OPC server data
|
||||
|
||||
Returns:
|
||||
None: Workflow completes without return value
|
||||
|
||||
Raises:
|
||||
Exception: If orchestration operations fail
|
||||
"""
|
||||
|
||||
input_data['workflow_name'] = 'orchestrator'
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data.get('schedule_name', 'orchestrator'),
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
'workflow_name': input_data['workflow_name'],
|
||||
}
|
||||
}
|
||||
|
||||
pipeline_config_handler = workflow.start_local_activity_method(
|
||||
Activities.aggregate_documents_in_mongodb,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['pipelines_query'],
|
||||
'timestamp_fields': ['updated_at'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
opc_servers_handler = workflow.start_local_activity_method(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{**metadata, 'query': input_data['opc_servers_query']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
orchestrated_schedules_handler = workflow.start_local_activity_method(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{
|
||||
**metadata,
|
||||
'query': {'collection': 'orchestrated_schedules'},
|
||||
'timestamp_fields': ['updated_at'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
current_slot_config_handler = workflow.start_local_activity_method(
|
||||
Activities.load_opc_slots,
|
||||
{
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
active_ingestors_handler = workflow.start_local_activity_method(
|
||||
Activities.load_active_ingestors,
|
||||
{
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
pipeline_config = await pipeline_config_handler
|
||||
orchestrated_schedules = await orchestrated_schedules_handler
|
||||
current_slot_config = await current_slot_config_handler
|
||||
opc_servers = await opc_servers_handler
|
||||
active_ingestors = await active_ingestors_handler
|
||||
|
||||
formatted_orchestrated_schedules_handler = workflow.start_local_activity_method(
|
||||
Activities.format_schedule_config,
|
||||
{**metadata, 'schedule_config': orchestrated_schedules},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedules_config_handler = workflow.start_local_activity_method(
|
||||
Activities.process_schedules,
|
||||
{**metadata, 'pipelines': pipeline_config},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
slot_config_handler = workflow.start_local_activity_method(
|
||||
Activities.process_slots,
|
||||
{
|
||||
**metadata,
|
||||
'opc_servers': opc_servers,
|
||||
'active_ingestors': active_ingestors,
|
||||
'pipelines': pipeline_config,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedules_config = await schedules_config_handler
|
||||
slot_config = await slot_config_handler
|
||||
formatted_orchestrated_schedules = await formatted_orchestrated_schedules_handler
|
||||
|
||||
schedule_actions_handler = workflow.start_local_activity_method(
|
||||
Activities.create_schedule_config,
|
||||
{
|
||||
**metadata,
|
||||
'current_schedule_config': formatted_orchestrated_schedules,
|
||||
'schedule_config': schedules_config,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
slot_actions_handler = workflow.start_local_activity_method(
|
||||
Activities.create_slot_config,
|
||||
{**metadata, 'current_slot_config': current_slot_config, 'slot_config': slot_config},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
normalize_schedules_handler = workflow.start_activity_method(
|
||||
Activities.normalize_schedules,
|
||||
{**metadata, 'orchestrated_schedules': formatted_orchestrated_schedules},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
create_collection_with_ttl_index_handler = workflow.start_activity_method(
|
||||
Activities.create_collection_with_ttl_index,
|
||||
{**metadata, 'pipelines': schedules_config['scouter']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedule_actions = await schedule_actions_handler
|
||||
slot_actions = await slot_actions_handler
|
||||
await normalize_schedules_handler
|
||||
await create_collection_with_ttl_index_handler
|
||||
|
||||
slot_deletion_report_handler = workflow.start_activity_method(
|
||||
Activities.delete_slots,
|
||||
{**metadata, 'to_delete': slot_actions['to_delete']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
slot_insertion_report_handler = workflow.start_activity_method(
|
||||
Activities.update_slots,
|
||||
{**metadata, 'to_insert': slot_actions['to_insert']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedule_deletion_report_handler = workflow.start_activity_method(
|
||||
Activities.delete_schedules,
|
||||
{**metadata, 'schedules': schedule_actions['to_delete']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedule_insertion_report_handler = workflow.start_activity_method(
|
||||
Activities.create_schedules,
|
||||
{**metadata, 'schedules': schedule_actions['to_create']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
schedule_update_report_handler = workflow.start_activity_method(
|
||||
Activities.update_schedules,
|
||||
{**metadata, 'schedules': schedule_actions['to_update']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
slot_deletion_report = await slot_deletion_report_handler
|
||||
slot_insertion_report = await slot_insertion_report_handler
|
||||
schedule_deletion_report = await schedule_deletion_report_handler
|
||||
schedule_insertion_report = await schedule_insertion_report_handler
|
||||
schedule_update_report = await schedule_update_report_handler
|
||||
|
||||
if schedule_insertion_report or schedule_update_report or schedule_deletion_report:
|
||||
schedule_report_handler = workflow.start_activity_method(
|
||||
Activities.report_schedule_orchestration,
|
||||
{
|
||||
**metadata,
|
||||
'created_schedules': schedule_insertion_report,
|
||||
'updated_schedules': schedule_update_report,
|
||||
'deleted_schedules': schedule_deletion_report,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
if slot_insertion_report or slot_deletion_report:
|
||||
slot_report_handler = workflow.start_activity_method(
|
||||
Activities.report_slot_orchestration,
|
||||
{
|
||||
**metadata,
|
||||
'inserted_slots': slot_insertion_report,
|
||||
'deleted_slots': slot_deletion_report,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
if schedule_update_report:
|
||||
update_pipelines_timestamps_handler = workflow.start_activity_method(
|
||||
Activities.update_pipelines_timestamps,
|
||||
{**metadata, 'updated_pipelines': schedule_update_report},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
if schedule_insertion_report:
|
||||
create_pipelines_timestamps_handler = workflow.start_activity_method(
|
||||
Activities.create_pipelines_timestamps,
|
||||
{**metadata, 'created_pipelines': schedule_insertion_report},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
if schedule_deletion_report:
|
||||
delete_pipelines_timestamps_handler = workflow.start_activity_method(
|
||||
Activities.delete_pipelines_timestamps,
|
||||
{**metadata, 'deleted_pipelines': schedule_deletion_report},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
)
|
||||
|
||||
if schedule_insertion_report or schedule_update_report or schedule_deletion_report:
|
||||
await schedule_report_handler
|
||||
|
||||
if slot_insertion_report or slot_deletion_report:
|
||||
await slot_report_handler
|
||||
|
||||
if schedule_update_report:
|
||||
await update_pipelines_timestamps_handler
|
||||
|
||||
if schedule_insertion_report:
|
||||
await create_pipelines_timestamps_handler
|
||||
|
||||
if schedule_deletion_report:
|
||||
await delete_pipelines_timestamps_handler
|
||||
96
orchestrator/workflows/reports.py
Normal file
96
orchestrator/workflows/reports.py
Normal file
@@ -0,0 +1,96 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='reports')
|
||||
class Reports:
|
||||
"""
|
||||
Reports workflow for sending scheduled notification summaries.
|
||||
|
||||
This workflow processes and sends scheduled reports to configured
|
||||
user groups. It loads notification data from MongoDB, filters it
|
||||
by receiver group configurations, and sends formatted HTML reports
|
||||
via email.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Execute the reports workflow for scheduled notification delivery.
|
||||
|
||||
This workflow loads all notifications from the notification queue,
|
||||
applies receiver group filtering, and sends comprehensive HTML
|
||||
reports to configured user groups.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Workflow input parameters.
|
||||
Required fields:
|
||||
- schedule_name (str): Name of the report schedule
|
||||
|
||||
Returns:
|
||||
None: Workflow completes without return value
|
||||
|
||||
Raises:
|
||||
Exception: If report generation or delivery fails
|
||||
"""
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': input_data['schedule_name'],
|
||||
'workflow_name': 'reports',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
mail_type = 'Reports'
|
||||
|
||||
input_data['metadata'] = metadata
|
||||
input_data['mail_type'] = mail_type
|
||||
|
||||
input_data['base_data_filter'] = {}
|
||||
|
||||
# Call subworkflow "load_notification_package" passing the static filters
|
||||
# (timestamp > last timestamp)
|
||||
|
||||
package = await workflow.execute_child_workflow(
|
||||
'subworkflow.load_notification_package', input_data
|
||||
)
|
||||
|
||||
if not package['notification_package'] or not package['sending_configs']:
|
||||
return
|
||||
|
||||
# Filter notification package by groups custom configs, levels and
|
||||
# timestamp cached
|
||||
|
||||
receiver_groups = await workflow.execute_local_activity_method(
|
||||
Activities.filter_notification_reports,
|
||||
{
|
||||
**metadata,
|
||||
'notification_package': package['notification_package'],
|
||||
'sending_configs': package['sending_configs'],
|
||||
},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
if not receiver_groups:
|
||||
return
|
||||
|
||||
# Call subworkflow "process_notifications" passing the notification package
|
||||
await workflow.execute_child_workflow(
|
||||
'subworkflow.process_notifications',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'mail_type': mail_type,
|
||||
'notification_package': receiver_groups,
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_report',
|
||||
},
|
||||
)
|
||||
0
orchestrator/workflows/subworkflows/__init__.py
Normal file
0
orchestrator/workflows/subworkflows/__init__.py
Normal file
108
orchestrator/workflows/subworkflows/load_notification_package.py
Normal file
108
orchestrator/workflows/subworkflows/load_notification_package.py
Normal file
@@ -0,0 +1,108 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='subworkflow.load_notification_package')
|
||||
class LoadNotificationPackage:
|
||||
"""
|
||||
Subworkflow for loading notification data and configuration.
|
||||
|
||||
This subworkflow retrieves notification packages from MongoDB and
|
||||
loads receiver group configurations. It handles timestamp-based
|
||||
filtering for incremental data processing and manages the data
|
||||
required for notification workflows.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Load notification package and sending configurations.
|
||||
|
||||
This subworkflow loads notifications from the MongoDB notification
|
||||
queue using timestamp-based filtering and retrieves active receiver
|
||||
group configurations. It updates the last processed timestamp in Redis.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Workflow input parameters.
|
||||
Required fields:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- mail_type (str): Type of mail (Alerts/Reports)
|
||||
- base_data_filter (dict[str, Any]): Base filter for notification query
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Package containing:
|
||||
- last_timestamp (str | None): Last processed timestamp
|
||||
- notification_package (list[dict]): Retrieved notifications
|
||||
- sending_configs (list[dict]): Active receiver group configurations
|
||||
|
||||
Raises:
|
||||
Exception: If data loading fails
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
# Load last timestamp from redis "notification_last_timestamp"
|
||||
last_timestamp_handler = workflow.start_local_activity_method(
|
||||
Activities.get_last_data_timestamp,
|
||||
{**metadata, 'mail_type': input_data['mail_type']},
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
# In parallel, load sending configs from collection "receiver_groups"
|
||||
sending_configs_handler = workflow.start_local_activity_method(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{**metadata, 'query': {'collection': 'receiver_groups', 'filters': {'active': True}}},
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
last_timestamp = await last_timestamp_handler
|
||||
|
||||
# Load notification package from collection "notification_queue", using a
|
||||
# static filter
|
||||
|
||||
notification_package = await workflow.start_local_activity_method(
|
||||
Activities.load_latest_data,
|
||||
{
|
||||
**metadata,
|
||||
'collection_name': 'notification_queue',
|
||||
'last_data_timestamp': last_timestamp,
|
||||
'base_data_filter': input_data['base_data_filter'],
|
||||
},
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
sending_configs = await sending_configs_handler
|
||||
|
||||
if not sending_configs or not notification_package:
|
||||
return {
|
||||
'last_timestamp': last_timestamp,
|
||||
'notification_package': notification_package,
|
||||
'sending_configs': sending_configs,
|
||||
}
|
||||
|
||||
# Put last collected timestamp in redis "notification_last_timestamp"
|
||||
|
||||
await workflow.start_activity_method(
|
||||
Activities.put_last_data_timestamp,
|
||||
{**metadata, 'data': notification_package, 'mail_type': input_data['mail_type']},
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
# Return a dict with the following keys:
|
||||
# - last_timestamp
|
||||
# - notification_package
|
||||
# - sending_configs
|
||||
return {
|
||||
'last_timestamp': last_timestamp,
|
||||
'notification_package': notification_package,
|
||||
'sending_configs': sending_configs,
|
||||
}
|
||||
99
orchestrator/workflows/subworkflows/process_notifications.py
Normal file
99
orchestrator/workflows/subworkflows/process_notifications.py
Normal file
@@ -0,0 +1,99 @@
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
|
||||
|
||||
@workflow.defn(name='subworkflow.process_notifications')
|
||||
class ProcessNotifications:
|
||||
"""
|
||||
Subworkflow for processing and sending notification emails.
|
||||
|
||||
This subworkflow handles the email delivery process including HTML
|
||||
generation, email sending, and logging to PostgreSQL. It processes
|
||||
receiver groups and generates delivery reports for monitoring.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Process notifications and send emails to configured receiver groups.
|
||||
|
||||
This subworkflow builds HTML email content, sends emails to all
|
||||
receiver groups, and logs the delivery results to PostgreSQL for
|
||||
monitoring and audit purposes.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Workflow input parameters.
|
||||
Required fields:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- mail_type (str): Type of email being sent
|
||||
- schema (str): PostgreSQL schema name for logging
|
||||
- table_name (str): PostgreSQL table name for logging
|
||||
- notification_package (dict[str, Any]): Receiver groups with notifications
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Log report of email delivery results
|
||||
|
||||
Raises:
|
||||
Exception: If notification processing fails
|
||||
"""
|
||||
|
||||
metadata = input_data['metadata']
|
||||
|
||||
# Use notification package to create the report html for each group and each model
|
||||
data_to_sent = await workflow.execute_local_activity_method(
|
||||
Activities.build_email_html,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': input_data['notification_package'],
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
# Send the report html to the receivers of each group
|
||||
log_report = await workflow.execute_activity_method(
|
||||
Activities.send_email,
|
||||
{**metadata, 'receiver_groups': data_to_sent, 'mail_type': input_data['mail_type']},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
if not log_report:
|
||||
return {}
|
||||
|
||||
# Format the log report to a dataframe to be stored in the database
|
||||
log_report = await workflow.execute_local_activity_method(
|
||||
Activities.format_log_report,
|
||||
{**metadata, 'receiver_groups': log_report, 'mail_type': input_data['mail_type']},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
# Store sending log in postgres database "log_report"
|
||||
await workflow.execute_activity_method(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'data': log_report,
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
schedule_to_close_timeout=timedelta(seconds=60),
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
# Return the log report to the caller
|
||||
return log_report
|
||||
43
pipelines_sample.json
Normal file
43
pipelines_sample.json
Normal file
@@ -0,0 +1,43 @@
|
||||
{
|
||||
"schedule_name": "scouter-test_linreg",
|
||||
"workflow_type": "scouter",
|
||||
"active": true,
|
||||
"model_id": "1",
|
||||
"frequency": "1m",
|
||||
"offset": "0m",
|
||||
"max_retry_policy": 1,
|
||||
"execution_timeout_seconds": 300,
|
||||
"task_timeout_seconds": 300,
|
||||
"on_conflict": "error",
|
||||
"debug_data_package": false,
|
||||
"fill_missing_tags": false,
|
||||
"filters": [
|
||||
{ "filter_name": "OUT_OF_RANGE", "policy": "STOP" }
|
||||
],
|
||||
"read_tags": [
|
||||
{
|
||||
"server_id": "1",
|
||||
"tag_name": "Counter",
|
||||
"tag_address": "ns=2;s=GatewayProSYS.ProSYS.Simulation.Counter",
|
||||
"aggr_func": "lts",
|
||||
"data_range": [1, 30],
|
||||
"frequency": 30000
|
||||
},
|
||||
{
|
||||
"server_id": "1",
|
||||
"tag_name": "Square",
|
||||
"tag_address": "ns=2;s=GatewayProSYS.ProSYS.Simulation.Square",
|
||||
"aggr_func": "lts",
|
||||
"data_range": [-2, 2],
|
||||
"frequency": 30000
|
||||
},
|
||||
{
|
||||
"server_id": "1",
|
||||
"tag_name": "Rollout",
|
||||
"tag_address": "ns=2;s=GatewayProSYS.ProSYS.Simulation.Rollout",
|
||||
"aggr_func": "lts",
|
||||
"data_range": [-2, 2],
|
||||
"frequency": 30000
|
||||
}
|
||||
]
|
||||
}
|
||||
162
pyproject.toml
Normal file
162
pyproject.toml
Normal file
@@ -0,0 +1,162 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "orchestrator"
|
||||
version = "0.0.0"
|
||||
description = "Sientia DataOps Orchestrator - 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 = [
|
||||
"B023", # ignore blind assignment, we need to assign the schedule to the schedule action
|
||||
"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
|
||||
]
|
||||
"e2e/**/*.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"
|
||||
]
|
||||
markers = [
|
||||
"asyncio: marks tests as async",
|
||||
"integration: marks tests as integration tests",
|
||||
"unit: marks tests as unit tests",
|
||||
"e2e: end-to-end tests requiring Docker (testcontainers + Temporal local server)",
|
||||
]
|
||||
|
||||
[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
|
||||
24
requirements-dev.txt
Normal file
24
requirements-dev.txt
Normal file
@@ -0,0 +1,24 @@
|
||||
# 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)
|
||||
testcontainers[postgres,mongodb]>=4.0.0
|
||||
aiosmtpd>=1.4.0
|
||||
sqlalchemy>=2.0.0
|
||||
pymongo>=4.6.0
|
||||
redis>=5.0.0
|
||||
|
||||
# Development Tools
|
||||
ipython>=8.12.0 # Enhanced Python shell
|
||||
ipdb>=0.13.13 # IPython debugger
|
||||
8
requirements-local.txt
Normal file
8
requirements-local.txt
Normal file
@@ -0,0 +1,8 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy
|
||||
redis
|
||||
pymongo
|
||||
jinja2
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1
|
||||
prometheus-client
|
||||
8
requirements.txt
Normal file
8
requirements.txt
Normal file
@@ -0,0 +1,8 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy
|
||||
redis
|
||||
pymongo
|
||||
jinja2
|
||||
sientia_do>=1.12.1
|
||||
prometheus-client
|
||||
11
run_coverage.sh
Executable file
11
run_coverage.sh
Executable file
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Exit on any error
|
||||
set -e
|
||||
|
||||
echo "Activating virtual environment..."
|
||||
source ./venv/bin/activate
|
||||
|
||||
pytest --cov=orchestrator --cov-report=html
|
||||
|
||||
xdg-open htmlcov/index.html
|
||||
18
run_local.sh
Executable file
18
run_local.sh
Executable file
@@ -0,0 +1,18 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Exit on any error
|
||||
set -e
|
||||
|
||||
echo "Activating virtual environment..."
|
||||
source ./venv/bin/activate
|
||||
|
||||
echo "Loading environment variables from .env..."
|
||||
if [ -f .env ]; then
|
||||
export $(cat .env | grep -v '^#' | xargs)
|
||||
echo "Environment variables loaded from .env"
|
||||
else
|
||||
echo "Warning: .env file not found. Continuing without environment variables."
|
||||
fi
|
||||
|
||||
echo "Starting orchestrator application..."
|
||||
exec python -m orchestrator.worker.worker
|
||||
11
sonar-project.properties
Normal file
11
sonar-project.properties
Normal file
@@ -0,0 +1,11 @@
|
||||
sonar.projectKey=Aignosi_sientia-dataops-orchestrator_temporal_22f04ada-022d-4345-a74f-bbda07dc17a6
|
||||
sonar.projectName=sientia-dataops-orchestrator_temporal
|
||||
sonar.sources=orchestrator
|
||||
sonar.tests=tests
|
||||
sonar.qualitygate.wait=true
|
||||
sonar.qualitygate.timeout=300
|
||||
sonar.python.coverage.reportPaths=coverage.xml
|
||||
sonar.python.xunit.reportPath=pytest.xml
|
||||
sonar.python.version=3.11
|
||||
sonar.projectVersion=1.0.0
|
||||
sonar.coverage.exclusions=orchestrator/worker/*
|
||||
2686
test.ipynb
Normal file
2686
test.ipynb
Normal file
File diff suppressed because it is too large
Load Diff
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
0
tests/orchestrator/__init__.py
Normal file
0
tests/orchestrator/__init__.py
Normal file
0
tests/orchestrator/activities/__init__.py
Normal file
0
tests/orchestrator/activities/__init__.py
Normal file
148
tests/orchestrator/activities/test_activities.py
Normal file
148
tests/orchestrator/activities/test_activities.py
Normal file
@@ -0,0 +1,148 @@
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
from orchestrator.activities.mongo_db import MongoDB
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.MongoDB.__init__')
|
||||
@patch('orchestrator.activities.temporal_manager.TemporalManager.__init__')
|
||||
@patch('orchestrator.activities.slot_manager.SlotManager.__init__')
|
||||
@patch('orchestrator.activities.formatters.Formatters.__init__')
|
||||
@patch('orchestrator.activities.email.Email.__init__')
|
||||
@patch('sientia_do.temporal.activities.postgres_sync.Postgres.__init__')
|
||||
@patch('orchestrator.activities.activities.MetricsController')
|
||||
def test___init__(
|
||||
mock_metrics_controller,
|
||||
mock_postgres_init,
|
||||
mock_email_init,
|
||||
mock_formatters_init,
|
||||
mock_slot_manager_init,
|
||||
mock_temporal_manager_init,
|
||||
mock_mongodb_init,
|
||||
):
|
||||
mongo_db_config = {
|
||||
'connection_string': 'mongodb://localhost:27017',
|
||||
'database_name': 'test_db',
|
||||
'ttl_index_seconds': 3600,
|
||||
}
|
||||
|
||||
redis_config = {'host': 'localhost', 'port': 6379, 'username': 'admin', 'password': 'password'}
|
||||
|
||||
temporal_config = {
|
||||
'temporal_host': 'localhost',
|
||||
'temporal_scouter_namespace': 'scouter',
|
||||
'temporal_laborious_namespace': 'laborious',
|
||||
}
|
||||
|
||||
email_config = {
|
||||
'sender_email': 'test@test.com',
|
||||
'sender_password': 'test',
|
||||
'smtp_server': 'test',
|
||||
'smtp_port': 587,
|
||||
}
|
||||
|
||||
postgres_config = {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'admin',
|
||||
'password': 'password',
|
||||
'dbname': 'test_db',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
|
||||
activities = Activities(
|
||||
temporal_config=temporal_config,
|
||||
redis_config=redis_config,
|
||||
mongodb_config=mongo_db_config,
|
||||
email_config=email_config,
|
||||
postgres_config=postgres_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
)
|
||||
|
||||
assert isinstance(activities, Activities)
|
||||
assert isinstance(activities, MongoDB)
|
||||
assert isinstance(activities, TemporalManager)
|
||||
assert isinstance(activities, SlotManager)
|
||||
assert isinstance(activities, Formatters)
|
||||
|
||||
mock_slot_manager_init.assert_called_once_with(
|
||||
ANY,
|
||||
host='localhost',
|
||||
port=6379,
|
||||
username='admin',
|
||||
password='password',
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_mongodb_init.assert_called_once_with(
|
||||
ANY,
|
||||
connection_string='mongodb://localhost:27017',
|
||||
database_name='test_db',
|
||||
ttl_index_seconds=3600,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_temporal_manager_init.assert_called_once_with(
|
||||
ANY,
|
||||
host='localhost',
|
||||
scouter_namespace='scouter',
|
||||
laborious_namespace='laborious',
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
mock_formatters_init.assert_called_once_with(
|
||||
ANY,
|
||||
scouter_namespace='scouter',
|
||||
laborious_namespace='laborious',
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mock_metrics_controller.return_value,
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.activities.MongoDB')
|
||||
@patch('orchestrator.activities.activities.TemporalManager')
|
||||
@patch('orchestrator.activities.activities.SlotManager')
|
||||
@patch('orchestrator.activities.activities.Formatters')
|
||||
@patch('orchestrator.activities.activities.Email')
|
||||
@patch('orchestrator.activities.activities.Postgres')
|
||||
def test_shutdown(
|
||||
mock_mongodb,
|
||||
mock_temporal_manager,
|
||||
mock_slot_manager,
|
||||
mock_formatters,
|
||||
mock_email,
|
||||
mock_postgres,
|
||||
):
|
||||
activities = Activities(
|
||||
temporal_config=MagicMock(),
|
||||
redis_config=MagicMock(),
|
||||
mongodb_config=MagicMock(),
|
||||
email_config=MagicMock(),
|
||||
postgres_config=MagicMock(),
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
)
|
||||
|
||||
activities.shutdown()
|
||||
|
||||
mock_mongodb.close.assert_called()
|
||||
mock_temporal_manager.close.assert_called()
|
||||
mock_slot_manager.close.assert_called()
|
||||
mock_formatters.close.assert_called()
|
||||
mock_email.close.assert_called()
|
||||
mock_postgres.close.assert_called()
|
||||
326
tests/orchestrator/activities/test_email.py
Normal file
326
tests/orchestrator/activities/test_email.py
Normal file
@@ -0,0 +1,326 @@
|
||||
from smtplib import SMTPServerDisconnected
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
from orchestrator.activities.email import Email
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('orchestrator.activities.email.EmailBuilder')
|
||||
@patch('orchestrator.activities.email.smtplib')
|
||||
def email(smtplib, email_builder):
|
||||
email = Email(
|
||||
sender_email='test@test.com',
|
||||
sender_password='test',
|
||||
smtp_server='test',
|
||||
smtp_port=587,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
email.send_notification = MagicMock()
|
||||
email.emit_metric = AsyncMock()
|
||||
|
||||
return email
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.EmailBuilder')
|
||||
@patch('orchestrator.activities.email.smtplib')
|
||||
def test___init___with_password(smtplib, email_builder):
|
||||
email = Email(
|
||||
sender_email='test@test.com',
|
||||
sender_password='test',
|
||||
smtp_server='test',
|
||||
smtp_port=587,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
assert email.sender_email == 'test@test.com'
|
||||
assert email.sender_password == 'test'
|
||||
assert email.smtp_port == 587
|
||||
|
||||
smtplib.SMTP.assert_called_once_with('test', 587, timeout=20)
|
||||
smtplib.SMTP.return_value.starttls.assert_called_once()
|
||||
smtplib.SMTP.return_value.login.assert_called_once_with('test@test.com', 'test')
|
||||
|
||||
assert email.server == smtplib.SMTP.return_value
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.EmailBuilder')
|
||||
@patch('orchestrator.activities.email.smtplib')
|
||||
def test___init___without_password(smtplib, email_builder):
|
||||
email = Email(
|
||||
sender_email='test@test.com',
|
||||
sender_password=None,
|
||||
smtp_server='test',
|
||||
smtp_port=587,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
assert email.sender_email == 'test@test.com'
|
||||
assert email.sender_password is None
|
||||
assert email.smtp_port == 587
|
||||
|
||||
smtplib.SMTP.assert_called_once_with('test', 587, timeout=20)
|
||||
assert email.server == smtplib.SMTP.return_value
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test',
|
||||
'model_name': 'test',
|
||||
'model_id': 'test',
|
||||
'workflow_name': 'test',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.SientiaMonitoring')
|
||||
def test_close(sientia_monitoring_mock, email):
|
||||
email.close()
|
||||
|
||||
email.server.quit.assert_called_once()
|
||||
sientia_monitoring_mock.shutdown.assert_called_once()
|
||||
|
||||
|
||||
def test_build_email_html(email):
|
||||
email.email_builder.build_email = MagicMock(return_value='test')
|
||||
input_data = {
|
||||
**metadata,
|
||||
'receiver_groups': {
|
||||
'group_1': {
|
||||
'notifications': [
|
||||
{'type': 'test', 'subject': 'test', 'body': 'test'},
|
||||
{'type': 'test', 'subject': 'test', 'body': 'test'},
|
||||
]
|
||||
}
|
||||
},
|
||||
'mail_type': 'test',
|
||||
}
|
||||
|
||||
response = email.build_email_html(input_data)
|
||||
|
||||
assert response == {
|
||||
'group_1': {
|
||||
'notifications': [
|
||||
{
|
||||
'type': 'test',
|
||||
'subject': 'test',
|
||||
'body': 'test',
|
||||
},
|
||||
{
|
||||
'type': 'test',
|
||||
'subject': 'test',
|
||||
'body': 'test',
|
||||
},
|
||||
],
|
||||
'html': 'test',
|
||||
}
|
||||
}
|
||||
|
||||
email.email_builder.build_email.assert_called_once_with(
|
||||
input_data['receiver_groups']['group_1']['notifications'], input_data['mail_type']
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.MIMEBase')
|
||||
@patch('orchestrator.activities.email.encoders')
|
||||
def test_handle_attachments_success(encoders, mime_base, email):
|
||||
message = MagicMock()
|
||||
|
||||
attachments = [
|
||||
{'filename': 'file_1', 'attachment_content': 'test_content_1'},
|
||||
{'filename': 'file_2', 'attachment_content': 'test_content_2'},
|
||||
{'filename': 'file_3', 'attachment_content': 'test_content_3'},
|
||||
]
|
||||
|
||||
response = email.handle_attachments(attachments, message)
|
||||
|
||||
assert response == message
|
||||
|
||||
mime_base.assert_called_with('application', 'octet-stream')
|
||||
assert mime_base.return_value.set_payload.call_count == 3
|
||||
|
||||
mime_base.return_value.set_payload.assert_has_calls(
|
||||
[
|
||||
call(b'test_content_1'),
|
||||
call(b'test_content_2'),
|
||||
call(b'test_content_3'),
|
||||
]
|
||||
)
|
||||
|
||||
encoders.encode_base64.assert_called_with(mime_base.return_value)
|
||||
assert encoders.encode_base64.call_count == 3
|
||||
|
||||
mime_base.return_value.add_header.assert_has_calls(
|
||||
[
|
||||
call('Content-Disposition', 'attachment; filename="file_1"'),
|
||||
call('Content-Disposition', 'attachment; filename="file_2"'),
|
||||
call('Content-Disposition', 'attachment; filename="file_3"'),
|
||||
]
|
||||
)
|
||||
|
||||
message.attach.assert_called_with(mime_base.return_value)
|
||||
assert message.attach.call_count == 3
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.MIMEBase')
|
||||
def test_handle_attachments_failure(mime_base, email):
|
||||
message = MagicMock()
|
||||
|
||||
mime_base.side_effect = Exception('test')
|
||||
|
||||
attachments = [{'filename': 'file_1', 'attachment_content': 'test_content_1'}]
|
||||
|
||||
try:
|
||||
email.handle_attachments(attachments, message)
|
||||
except Exception as e:
|
||||
assert str(e) == 'test'
|
||||
|
||||
assert message.attach.call_count == 0
|
||||
|
||||
|
||||
def test_try_send_email_success(email):
|
||||
email.server.sendmail = MagicMock()
|
||||
|
||||
msg = MagicMock()
|
||||
|
||||
email.try_send_email(msg, 'test')
|
||||
|
||||
email.server.sendmail.assert_called_once_with(
|
||||
'test@test.com', 'test', msg.as_string.return_value
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.smtplib.SMTP')
|
||||
def test_try_send_email_reconnect_quit_success(smtp, email):
|
||||
email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
|
||||
email.server.quit = MagicMock()
|
||||
|
||||
msg = MagicMock()
|
||||
|
||||
email.try_send_email(msg, 'test')
|
||||
|
||||
smtp.assert_has_calls([call('test', 587, timeout=20)])
|
||||
|
||||
smtp.return_value.starttls.assert_called_once()
|
||||
smtp.return_value.login.assert_called_once_with('test@test.com', 'test')
|
||||
|
||||
smtp.return_value.sendmail.assert_called_once_with(
|
||||
'test@test.com', 'test', msg.as_string.return_value
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.smtplib.SMTP')
|
||||
def test_try_send_email_reconnect_quit_failure_disconnect(smtp, email):
|
||||
email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
|
||||
email.server.quit = MagicMock(side_effect=SMTPServerDisconnected('test'))
|
||||
|
||||
msg = MagicMock()
|
||||
|
||||
email.try_send_email(msg, 'test')
|
||||
|
||||
smtp.assert_has_calls([call('test', 587, timeout=20)])
|
||||
smtp.return_value.starttls.assert_called_once()
|
||||
smtp.return_value.login.assert_called_once_with('test@test.com', 'test')
|
||||
|
||||
smtp.return_value.sendmail.assert_called_once_with(
|
||||
'test@test.com', 'test', msg.as_string.return_value
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.smtplib.SMTP')
|
||||
def test_try_send_email_reconnect_quit_failure(smtp, email):
|
||||
email.server.sendmail = MagicMock(side_effect=SMTPServerDisconnected('test'))
|
||||
email.server.quit = MagicMock(side_effect=Exception('test'))
|
||||
|
||||
msg = MagicMock()
|
||||
|
||||
try:
|
||||
email.try_send_email(msg, 'test')
|
||||
except Exception as e:
|
||||
assert str(e) == 'test'
|
||||
else:
|
||||
raise AssertionError('Expected exception')
|
||||
|
||||
|
||||
def test_send_email_without_smtp_server(email):
|
||||
email.smtp_server = None
|
||||
input_data = {**metadata, 'receiver_groups': {}, 'mail_type': 'test_TYPE'}
|
||||
response = email.send_email(input_data)
|
||||
|
||||
assert response == {}
|
||||
|
||||
|
||||
@patch('orchestrator.activities.email.MIMEText')
|
||||
@patch('orchestrator.activities.email.MIMEMultipart')
|
||||
def test_send_email(mimemultipart, mimetext, email):
|
||||
side_effect_1 = MagicMock()
|
||||
side_effect_2 = MagicMock()
|
||||
mimemultipart.side_effect = [side_effect_1, side_effect_2]
|
||||
|
||||
email.email_builder.send_notification = MagicMock()
|
||||
|
||||
email.try_send_email = MagicMock(side_effect=[None, Exception('test')])
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'receiver_groups': {
|
||||
'group_1': {
|
||||
'members': ['test@test.com', 'test2@test.com'],
|
||||
'notifications': [
|
||||
{
|
||||
'attachment_content': 'test_content_1',
|
||||
'trigger': 'test_trigger',
|
||||
'notification_id': 'test_notification_id',
|
||||
}
|
||||
],
|
||||
'html': 'test_html1',
|
||||
},
|
||||
'group_2': {
|
||||
'members': ['test3@test.com', 'test4@test.com'],
|
||||
'notifications': [],
|
||||
'html': 'test_html2',
|
||||
},
|
||||
},
|
||||
'mail_type': 'test_TYPE',
|
||||
}
|
||||
|
||||
response = email.send_email(input_data)
|
||||
|
||||
assert response['group_1']['status'] == 'sent'
|
||||
assert response['group_2']['status'] == 'failed'
|
||||
|
||||
assert mimemultipart.call_count == 2
|
||||
|
||||
mimetext.assert_has_calls([call('test_html1', 'html'), call('test_html2', 'html')])
|
||||
|
||||
side_effect_1.__setitem__.assert_has_calls(
|
||||
[
|
||||
call('From', 'test@test.com'),
|
||||
call('To', 'test@test.com, test2@test.com'),
|
||||
call('Subject', 'SIENTIA™ test_TYPE'),
|
||||
]
|
||||
)
|
||||
|
||||
side_effect_2.__setitem__.assert_has_calls(
|
||||
[
|
||||
call('From', 'test@test.com'),
|
||||
call('To', 'test3@test.com, test4@test.com'),
|
||||
call('Subject', 'SIENTIA™ test_TYPE'),
|
||||
]
|
||||
)
|
||||
|
||||
email.try_send_email.assert_has_calls(
|
||||
[
|
||||
call(side_effect_1, 'test@test.com, test2@test.com'),
|
||||
call(side_effect_2, 'test3@test.com, test4@test.com'),
|
||||
]
|
||||
)
|
||||
|
||||
assert email.try_send_email.call_count == 2
|
||||
768
tests/orchestrator/activities/test_formatters.py
Normal file
768
tests/orchestrator/activities/test_formatters.py
Normal file
@@ -0,0 +1,768 @@
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pandas import DataFrame
|
||||
from pytest import fixture
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
|
||||
|
||||
@fixture
|
||||
def formatters():
|
||||
formatters = Formatters(
|
||||
scouter_namespace='scouter',
|
||||
laborious_namespace='laborious',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
formatters.send_notification = MagicMock()
|
||||
formatters.emit_metric = AsyncMock()
|
||||
formatters.error = MagicMock()
|
||||
formatters.info = MagicMock()
|
||||
formatters.debug = MagicMock()
|
||||
return formatters
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test_schedule_name',
|
||||
'workflow_name': 'test_workflow_name',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_process_schedules(formatters):
|
||||
mock_scouter = MagicMock(return_value={'test_scouter': 'test_scouter'})
|
||||
mock_predictions_batch = MagicMock(
|
||||
return_value={'test_predictions_batch': 'test_predictions_batch'}
|
||||
)
|
||||
mock_minimal_retrain = MagicMock(return_value={'test_minimal_retrain': 'test_minimal_retrain'})
|
||||
mock_drift = MagicMock(return_value={'test_drift': 'test_drift'})
|
||||
mock_simple_metrics = MagicMock(return_value={'test_simple_metrics': 'test_simple_metrics'})
|
||||
|
||||
mock_schedule_types = {
|
||||
'scouter': {
|
||||
'namespace': 'scouter',
|
||||
'function': mock_scouter,
|
||||
},
|
||||
'pi_web_api_scouter': {
|
||||
'namespace': 'scouter',
|
||||
'function': mock_scouter,
|
||||
},
|
||||
'predictions_batch': {
|
||||
'namespace': 'laborious',
|
||||
'function': mock_predictions_batch,
|
||||
},
|
||||
'minimal_retrain': {
|
||||
'namespace': 'laborious',
|
||||
'function': mock_minimal_retrain,
|
||||
},
|
||||
'drift': {
|
||||
'namespace': 'laborious',
|
||||
'function': mock_drift,
|
||||
},
|
||||
'simple_metrics': {
|
||||
'namespace': 'laborious',
|
||||
'function': mock_simple_metrics,
|
||||
},
|
||||
}
|
||||
|
||||
input_data = {
|
||||
'pipelines': [
|
||||
{
|
||||
'schedule_name': 'test_schedule_name',
|
||||
'workflow_type': 'scouter',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-01',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name2',
|
||||
'workflow_type': 'predictions_batch',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-02',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name3',
|
||||
'workflow_type': 'minimal_retrain',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-03',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name4',
|
||||
'workflow_type': 'drift',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-04',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name5',
|
||||
'workflow_type': 'simple_metrics',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-05',
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
with patch('orchestrator.activities.formatters.schedule_types', mock_schedule_types):
|
||||
result = formatters.process_schedules(input_data)
|
||||
|
||||
assert result == {
|
||||
'scouter': {
|
||||
'test_schedule_name': {'test_scouter': 'test_scouter', 'updated_at': '2021-01-01'}
|
||||
},
|
||||
'laborious': {
|
||||
'test_schedule_name2': {
|
||||
'test_predictions_batch': 'test_predictions_batch',
|
||||
'updated_at': '2021-01-02',
|
||||
},
|
||||
'test_schedule_name3': {
|
||||
'test_minimal_retrain': 'test_minimal_retrain',
|
||||
'updated_at': '2021-01-03',
|
||||
},
|
||||
'test_schedule_name4': {
|
||||
'test_drift': 'test_drift',
|
||||
'updated_at': '2021-01-04',
|
||||
},
|
||||
'test_schedule_name5': {
|
||||
'test_simple_metrics': 'test_simple_metrics',
|
||||
'updated_at': '2021-01-05',
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
mock_scouter.assert_called_once_with(input_data['pipelines'][0])
|
||||
mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1])
|
||||
mock_minimal_retrain.assert_called_once_with(input_data['pipelines'][2])
|
||||
mock_drift.assert_called_once_with(input_data['pipelines'][3])
|
||||
mock_simple_metrics.assert_called_once_with(input_data['pipelines'][4])
|
||||
|
||||
|
||||
def test_process_schedules_with_invalid_workflow_type(formatters):
|
||||
"""Test that process_schedules handles invalid workflow types correctly"""
|
||||
mock_scouter = MagicMock(return_value={'test_scouter': 'test_scouter'})
|
||||
|
||||
mock_schedule_types = {
|
||||
'scouter': {
|
||||
'namespace': 'scouter',
|
||||
'function': mock_scouter,
|
||||
},
|
||||
}
|
||||
|
||||
pipelines = [
|
||||
{
|
||||
'schedule_name': 'test_schedule_name_valid',
|
||||
'workflow_type': 'scouter',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-01',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name_invalid',
|
||||
'workflow_type': 'invalid_workflow_type',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-02',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test_schedule_name_valid2',
|
||||
'workflow_type': 'scouter',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'updated_at': '2021-01-03',
|
||||
},
|
||||
]
|
||||
|
||||
input_data = {
|
||||
'pipelines': pipelines,
|
||||
'metadata': {
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
},
|
||||
}
|
||||
|
||||
with patch('orchestrator.activities.formatters.schedule_types', mock_schedule_types):
|
||||
result = formatters.process_schedules(input_data)
|
||||
|
||||
# Assert that error was called for invalid workflow type
|
||||
formatters.error.assert_called_once_with(
|
||||
'Workflow type invalid_workflow_type not supported', metadata=input_data['metadata']
|
||||
)
|
||||
|
||||
# Assert that only valid pipelines were processed
|
||||
assert result == {
|
||||
'scouter': {
|
||||
'test_schedule_name_valid': {
|
||||
'test_scouter': 'test_scouter',
|
||||
'updated_at': '2021-01-01',
|
||||
},
|
||||
'test_schedule_name_valid2': {
|
||||
'test_scouter': 'test_scouter',
|
||||
'updated_at': '2021-01-03',
|
||||
},
|
||||
},
|
||||
'laborious': {},
|
||||
}
|
||||
|
||||
# Assert that the mock function was called only for valid pipelines
|
||||
assert mock_scouter.call_count == 2
|
||||
mock_scouter.assert_any_call(pipelines[0])
|
||||
mock_scouter.assert_any_call(pipelines[2])
|
||||
|
||||
|
||||
@patch(
|
||||
'orchestrator.activities.formatters.gather_read_tags',
|
||||
return_value={
|
||||
'1:test_tag_address': {
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
'topics': ['raw_test_schedule'],
|
||||
'frequency': 1000,
|
||||
},
|
||||
'2:test_tag_address2': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address2',
|
||||
'topics': ['raw_test_schedule2'],
|
||||
'frequency': 1000,
|
||||
},
|
||||
'2:test_tag_address3': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address3',
|
||||
'topics': ['raw_test_schedule2'],
|
||||
'frequency': 1000,
|
||||
},
|
||||
},
|
||||
)
|
||||
@patch('orchestrator.activities.formatters.build_tag_config')
|
||||
def test_process_slots(mock_build_tag_config, mock_gather_read_tags, formatters):
|
||||
input_data = {
|
||||
'opc_servers': [
|
||||
{'id': '1', 'server_name': 'test_server_name', 'url': 'test_url', 'uri': 'test_uri'},
|
||||
{'id': '2', 'server_name': 'test_server_name2', 'url': 'test_url2', 'uri': 'test_uri2'},
|
||||
],
|
||||
'active_ingestors': ['test_active_ingestor1', 'test_active_ingestor2'],
|
||||
'pipelines': 'test_gather_read_tags',
|
||||
**metadata,
|
||||
}
|
||||
|
||||
expected_opc_servers = {
|
||||
'1': {
|
||||
'id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'url': 'test_url',
|
||||
'uri': 'test_uri',
|
||||
},
|
||||
'2': {
|
||||
'id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'url': 'test_url2',
|
||||
'uri': 'test_uri2',
|
||||
},
|
||||
}
|
||||
|
||||
slot_mock = {
|
||||
'test_server_name': {
|
||||
'server_id': '1',
|
||||
'name': 'test_server_name',
|
||||
'url': 'test_url',
|
||||
'server_uri': 'test_uri',
|
||||
'cert_path': None,
|
||||
'private_key_path': None,
|
||||
'server_cert_path': None,
|
||||
}
|
||||
}
|
||||
|
||||
mock_build_tag_config.return_value = (slot_mock, ['2'])
|
||||
|
||||
result = formatters.process_slots(input_data)
|
||||
|
||||
tags = list(mock_gather_read_tags.return_value.values())
|
||||
|
||||
mock_gather_read_tags.assert_called_once_with(input_data['pipelines'])
|
||||
|
||||
mock_build_tag_config.assert_has_calls(
|
||||
[
|
||||
call(tags[:2], expected_opc_servers),
|
||||
call(tags[2:], expected_opc_servers),
|
||||
]
|
||||
)
|
||||
|
||||
formatters.send_notification.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
|
||||
message='Servers 2 not found in opc_servers',
|
||||
block='orchestrator',
|
||||
level=NotificationLevel.ERROR,
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
|
||||
message='Servers 2 not found in opc_servers',
|
||||
block='orchestrator',
|
||||
level=NotificationLevel.ERROR,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
assert result == {
|
||||
'1': slot_mock,
|
||||
'2': slot_mock,
|
||||
}
|
||||
|
||||
|
||||
def test_format_schedule_config(formatters):
|
||||
input_data = {
|
||||
'schedule_config': [
|
||||
{'namespace': 'test_namespace1', 'schedule_name': 'test1', 'updated_at': '2021-01-01'},
|
||||
{'namespace': 'test_namespace2', 'schedule_name': 'test2', 'updated_at': '2021-01-02'},
|
||||
],
|
||||
**metadata,
|
||||
}
|
||||
|
||||
result = formatters.format_schedule_config(input_data)
|
||||
|
||||
assert result == {
|
||||
'test_namespace1': {'test1': '2021-01-01'},
|
||||
'test_namespace2': {'test2': '2021-01-02'},
|
||||
}
|
||||
|
||||
|
||||
def test_create_schedule_config(formatters):
|
||||
input_data = {
|
||||
'current_schedule_config': {
|
||||
'scouter': {
|
||||
'test_schedule_name_to_delete': '2021-01-01',
|
||||
'test_schedule_name_to_update': '2021-01-02',
|
||||
}
|
||||
},
|
||||
'schedule_config': {
|
||||
'laborious': {
|
||||
'test_schedule_name_to_create': {
|
||||
'frequency': 60,
|
||||
'data': {'test': 'test'},
|
||||
'updated_at': '2021-01-03',
|
||||
}
|
||||
},
|
||||
'scouter': {
|
||||
'test_schedule_name_to_update': {
|
||||
'frequency': 60,
|
||||
'data': {'test': 'test2'},
|
||||
'updated_at': '2021-01-04',
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
result = formatters.create_schedule_config(input_data)
|
||||
|
||||
assert result == {
|
||||
'to_create': {
|
||||
'laborious': {
|
||||
'test_schedule_name_to_create': {
|
||||
'frequency': 60,
|
||||
'data': {'test': 'test'},
|
||||
'updated_at': '2021-01-03',
|
||||
}
|
||||
},
|
||||
'scouter': {},
|
||||
},
|
||||
'to_update': {
|
||||
'scouter': {
|
||||
'test_schedule_name_to_update': {
|
||||
'frequency': 60,
|
||||
'data': {'test': 'test2'},
|
||||
'updated_at': '2021-01-04',
|
||||
}
|
||||
},
|
||||
'laborious': {},
|
||||
},
|
||||
'to_delete': {'scouter': ['test_schedule_name_to_delete'], 'laborious': []},
|
||||
}
|
||||
|
||||
|
||||
def test_create_slot_config(formatters):
|
||||
input_data = {
|
||||
'current_slot_config': {
|
||||
'1': {'frequency': 60, 'data': {'test': 'test'}},
|
||||
'2': {'frequency': 60, 'data': {'test': 'test'}},
|
||||
},
|
||||
'slot_config': {'1': {'frequency': 60, 'data': {'test': 'test2'}}},
|
||||
}
|
||||
|
||||
result = formatters.create_slot_config(input_data)
|
||||
|
||||
assert result == {
|
||||
'to_delete': ['2'],
|
||||
'to_insert': {'1': {'frequency': 60, 'data': {'test': 'test2'}}},
|
||||
}
|
||||
|
||||
|
||||
def test_send_success_report(formatters):
|
||||
formatters.send_success_report(
|
||||
metadata=metadata,
|
||||
message='test_message',
|
||||
notification_id='test_notification_id',
|
||||
attachment={'test': 'test'},
|
||||
)
|
||||
formatters.send_notification.assert_called_once_with(
|
||||
metadata=metadata,
|
||||
notification_id='test_notification_id',
|
||||
message='test_message',
|
||||
block='report_orchestration',
|
||||
level=NotificationLevel.INFO,
|
||||
attachment_content=json.dumps({'test': 'test'}, indent=4, sort_keys=True),
|
||||
)
|
||||
|
||||
|
||||
def test_send_error_report(formatters):
|
||||
formatters.send_error_report(
|
||||
metadata=metadata,
|
||||
message='test_message',
|
||||
notification_id='test_notification_id',
|
||||
attachment='test_attachment',
|
||||
)
|
||||
formatters.send_notification.assert_called_once_with(
|
||||
metadata=metadata,
|
||||
notification_id='test_notification_id',
|
||||
message='test_message',
|
||||
block='report_orchestration',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content='test_attachment',
|
||||
)
|
||||
|
||||
|
||||
def test_parse_report(formatters):
|
||||
input_data = {
|
||||
'test_key': {'success': True},
|
||||
'test_key2': {'success': False, 'message': 'test_error'},
|
||||
}
|
||||
|
||||
result = formatters.parse_report(input_data)
|
||||
|
||||
assert result == (['test_key'], ['test_key2'])
|
||||
|
||||
|
||||
def test_parse_report_schedule(formatters):
|
||||
input_data = [
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_create',
|
||||
'success': True,
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_create_error',
|
||||
'success': False,
|
||||
'message': 'test_error',
|
||||
'attachment': 'test_attachment',
|
||||
},
|
||||
]
|
||||
result = formatters.parse_report_schedule(input_data)
|
||||
|
||||
assert result == (
|
||||
['test_namespace/test_schedule_name_to_create'],
|
||||
{
|
||||
'test_namespace/test_schedule_name_to_create_error': {
|
||||
'message': 'test_error',
|
||||
'attachment': 'test_attachment',
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_report_schedule_orchestration(formatters):
|
||||
formatters.parse_report_schedule = MagicMock(side_effect=formatters.parse_report_schedule)
|
||||
formatters.send_success_report = AsyncMock()
|
||||
formatters.send_error_report = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'created_schedules': [
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_create',
|
||||
'success': True,
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_create_error',
|
||||
'success': False,
|
||||
'message': 'test_error1',
|
||||
'attachment': 'test_attachment1',
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_create_error2',
|
||||
'success': False,
|
||||
'message': 'test_error2',
|
||||
},
|
||||
],
|
||||
'updated_schedules': [
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_update',
|
||||
'success': True,
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_update_error',
|
||||
'success': False,
|
||||
'message': 'test_error2',
|
||||
'attachment': 'test_attachment2',
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_update_error2',
|
||||
'success': False,
|
||||
'message': 'test_error3',
|
||||
'attachment': 'test_attachment3',
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_update_error3',
|
||||
'success': False,
|
||||
'message': 'test_error4',
|
||||
},
|
||||
],
|
||||
'deleted_schedules': [
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_delete',
|
||||
'success': True,
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_delete_error',
|
||||
'success': False,
|
||||
'message': 'test_error4',
|
||||
'attachment': 'test_attachment4',
|
||||
},
|
||||
{
|
||||
'namespace': 'test_namespace',
|
||||
'schedule_name': 'test_schedule_name_to_delete_error2',
|
||||
'success': False,
|
||||
'message': 'test_error5',
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
formatters.report_schedule_orchestration(input_data)
|
||||
|
||||
formatters.parse_report_schedule.assert_has_calls(
|
||||
[
|
||||
call(input_data['created_schedules']),
|
||||
call(input_data['updated_schedules']),
|
||||
call(input_data['deleted_schedules']),
|
||||
]
|
||||
)
|
||||
formatters.send_success_report.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Successfully created schedules: \n test_namespace/test_schedule_name_to_create',
|
||||
notification_id='REPORT_ORCHESTRATION_CREATED_SCHEDULES',
|
||||
attachment=input_data['created_schedules'],
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Successfully updated schedules: \n test_namespace/test_schedule_name_to_update',
|
||||
notification_id='REPORT_ORCHESTRATION_UPDATED_SCHEDULES',
|
||||
attachment=input_data['updated_schedules'],
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Successfully deleted schedules: \n test_namespace/test_schedule_name_to_delete',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SCHEDULES',
|
||||
attachment=input_data['deleted_schedules'],
|
||||
),
|
||||
]
|
||||
)
|
||||
formatters.send_error_report.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Fails on created schedules: \n test_namespace/test_schedule_name_to_create_error, test_namespace/test_schedule_name_to_create_error2',
|
||||
notification_id='REPORT_ORCHESTRATION_CREATED_SCHEDULES_ERROR',
|
||||
attachment='test_namespace/test_schedule_name_to_create_error:\ntest_error1\ntest_attachment1\n ========== \ntest_namespace/test_schedule_name_to_create_error2:\ntest_error2',
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Fails on updated schedules: \n test_namespace/test_schedule_name_to_update_error, test_namespace/test_schedule_name_to_update_error2, test_namespace/test_schedule_name_to_update_error3',
|
||||
notification_id='REPORT_ORCHESTRATION_UPDATED_SCHEDULES_ERROR',
|
||||
attachment='test_namespace/test_schedule_name_to_update_error:\ntest_error2\ntest_attachment2\n ========== \ntest_namespace/test_schedule_name_to_update_error2:\ntest_error3\ntest_attachment3\n ========== \ntest_namespace/test_schedule_name_to_update_error3:\ntest_error4',
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Fails on deleted schedules: \n test_namespace/test_schedule_name_to_delete_error, test_namespace/test_schedule_name_to_delete_error2',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SCHEDULES_ERROR',
|
||||
attachment='test_namespace/test_schedule_name_to_delete_error:\ntest_error4\ntest_attachment4\n ========== \ntest_namespace/test_schedule_name_to_delete_error2:\ntest_error5',
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_report_slot_orchestration(formatters):
|
||||
formatters.parse_report = MagicMock(side_effect=formatters.parse_report)
|
||||
formatters.send_success_report = AsyncMock()
|
||||
formatters.send_error_report = AsyncMock()
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'inserted_slots': {
|
||||
'test_slot_name_to_create': {'success': True},
|
||||
'test_slot_name_to_create_error': {'success': False, 'error': 'test_error'},
|
||||
},
|
||||
'deleted_slots': {
|
||||
'test_slot_name_to_delete': {'success': True},
|
||||
'test_slot_name_to_delete_error': {'success': False, 'error': 'test_error'},
|
||||
},
|
||||
}
|
||||
|
||||
formatters.report_slot_orchestration(input_data)
|
||||
|
||||
formatters.parse_report.assert_has_calls(
|
||||
[call(input_data['inserted_slots']), call(input_data['deleted_slots'])]
|
||||
)
|
||||
formatters.send_success_report.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Inserted slots: \n test_slot_name_to_create',
|
||||
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Deleted slots: \n test_slot_name_to_delete',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
|
||||
),
|
||||
]
|
||||
)
|
||||
formatters.send_error_report.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Failed to insert slots: \n test_slot_name_to_create_error',
|
||||
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
|
||||
attachment=input_data['inserted_slots'],
|
||||
),
|
||||
call(
|
||||
metadata=metadata['metadata'],
|
||||
message='Failed to delete slots: \n test_slot_name_to_delete_error',
|
||||
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
|
||||
attachment=input_data['deleted_slots'],
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_format_log_report(formatters):
|
||||
input_data = {
|
||||
**metadata,
|
||||
'receiver_groups': {
|
||||
'test_receiver_group': {
|
||||
'notifications': [
|
||||
{
|
||||
'notification_id': 'test_notification_id',
|
||||
'trigger': 'test_trigger',
|
||||
'timestamp': '2021-01-01',
|
||||
'message': 'test_message',
|
||||
'level': 'test_level',
|
||||
'block': 'test_block',
|
||||
'pipeline': 'test_pipeline',
|
||||
'project': 'test_project',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
],
|
||||
'status': 'sent',
|
||||
},
|
||||
'test_receiver_group2': {
|
||||
'notifications': [
|
||||
{
|
||||
'notification_id': 'test_notification_id',
|
||||
'trigger': 'test_trigger',
|
||||
'timestamp': '2021-01-01',
|
||||
'message': 'test_message',
|
||||
'level': 'test_level',
|
||||
'block': 'test_block',
|
||||
'pipeline': 'test_pipeline',
|
||||
'project': 'test_project',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
],
|
||||
'status': 'sent',
|
||||
},
|
||||
},
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
result = formatters.format_log_report(input_data)
|
||||
|
||||
expected_result = DataFrame(
|
||||
[
|
||||
{
|
||||
'status': 'sent',
|
||||
'timestamp': '2021-01-01',
|
||||
'groups': ['test_receiver_group', 'test_receiver_group2'],
|
||||
'message': 'test_message',
|
||||
'level': 'test_level',
|
||||
'notification_id': 'test_notification_id',
|
||||
'block': 'test_block',
|
||||
'schedule': 'test_trigger',
|
||||
'pipeline': 'test_pipeline',
|
||||
'project': 'test_project',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert DataFrame(result).equals(expected_result)
|
||||
|
||||
|
||||
def test_filter_notification_reports(formatters):
|
||||
input_data = {
|
||||
**metadata,
|
||||
'notification_package': [
|
||||
{'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'},
|
||||
{'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
|
||||
{'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'},
|
||||
],
|
||||
'sending_configs': [
|
||||
{
|
||||
'group_name': 'test_group_1',
|
||||
'contents': ['reports'],
|
||||
'ignore': ['test_notification_id_1'],
|
||||
},
|
||||
{'group_name': 'test_group_2', 'contents': ['core_alerts']},
|
||||
],
|
||||
}
|
||||
|
||||
response = formatters.filter_notification_reports(input_data)
|
||||
|
||||
assert response == {
|
||||
'test_group_1': {
|
||||
'group_name': 'test_group_1',
|
||||
'contents': ['reports'],
|
||||
'ignore': ['test_notification_id_1'],
|
||||
'notifications': [
|
||||
{'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
|
||||
{'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'},
|
||||
],
|
||||
}
|
||||
}
|
||||
545
tests/orchestrator/activities/test_mongo_db.py
Normal file
545
tests/orchestrator/activities/test_mongo_db.py
Normal file
@@ -0,0 +1,545 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import fixture
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
|
||||
|
||||
from orchestrator.activities.mongo_db import MongoDB
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('orchestrator.activities.mongo_db.MongoDBRepository')
|
||||
def mongo_db(mongo_mock):
|
||||
mongo = MongoDB(
|
||||
connection_string='mongodb://localhost:27017',
|
||||
database_name='test_db',
|
||||
ttl_index_seconds=3600,
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
mongo.send_notification = MagicMock()
|
||||
mongo.emit_metric = AsyncMock()
|
||||
|
||||
return mongo
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.MongoDBRepository')
|
||||
def test___init__(mongo_mock):
|
||||
logger = MagicMock()
|
||||
notification_handler = MagicMock()
|
||||
metrics_controller = AsyncMock()
|
||||
|
||||
MongoDB(
|
||||
connection_string='mongodb://localhost:27017',
|
||||
database_name='test_db',
|
||||
ttl_index_seconds=3600,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
mongo_mock.assert_called_once_with(
|
||||
connection_string='mongodb://localhost:27017',
|
||||
database_name='test_db',
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.SientiaMonitoring')
|
||||
def test_close(sientia_monitoring_mock, mongo_db):
|
||||
mongo_db.close()
|
||||
mongo_db.mongo_db_repository.close.assert_called_once()
|
||||
sientia_monitoring_mock.shutdown.assert_called_once()
|
||||
|
||||
|
||||
def test___del__(mongo_db):
|
||||
mongo_db.close = MagicMock()
|
||||
mongo_db.__del__()
|
||||
mongo_db.close.assert_called_once()
|
||||
|
||||
|
||||
def test_find_documents_in_mongodb_success(mongo_db):
|
||||
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
|
||||
|
||||
mongo_db.mongo_db_repository.find = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
'name': 'test1',
|
||||
'timestamp': datetime.strptime(
|
||||
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
|
||||
),
|
||||
},
|
||||
{
|
||||
'name': 'test2',
|
||||
'timestamp': datetime.strptime(
|
||||
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
|
||||
),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
result = mongo_db.find_documents_in_mongodb(
|
||||
{'query': input_data, 'timestamp_fields': ['timestamp']}
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0] == {'name': 'test1', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
|
||||
assert result[1] == {'name': 'test2', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
|
||||
mongo_db.mongo_db_repository.find.assert_called_once_with(
|
||||
'test_collection', {'name': {'$exists': True}}, {}
|
||||
)
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test_schedule_name',
|
||||
'workflow_name': 'test_workflow_name',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_find_documents_in_mongodb_failure(mongo_db):
|
||||
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
|
||||
mongo_db.mongo_db_repository.find = MagicMock(side_effect=Exception('Error'))
|
||||
|
||||
try:
|
||||
mongo_db.find_documents_in_mongodb({'query': input_data, **metadata})
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_QUERY_ERROR',
|
||||
message='Failed to execute MongoDB query: Error',
|
||||
level=NotificationLevel.ERROR,
|
||||
block='load_query_from_mongodb',
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_find_documents_in_mongodb_missing_collection(mongo_db):
|
||||
input_data = {'query': {'filters': {}}}
|
||||
|
||||
try:
|
||||
mongo_db.find_documents_in_mongodb(input_data)
|
||||
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Collection name must be provided in the query.'
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected a ValueError to be raised')
|
||||
|
||||
|
||||
def test_aggregate_documents_in_mongodb_success(mongo_db):
|
||||
input_data = {
|
||||
'collection': 'test_collection',
|
||||
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
|
||||
}
|
||||
mongo_db.mongo_db_repository.aggregate = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
'name': 'test1',
|
||||
'timestamp': datetime.strptime(
|
||||
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
|
||||
),
|
||||
},
|
||||
{
|
||||
'name': 'test2',
|
||||
'timestamp': datetime.strptime(
|
||||
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
|
||||
),
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
result = mongo_db.aggregate_documents_in_mongodb(
|
||||
{'query': input_data, 'timestamp_fields': ['timestamp']}
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0] == {'name': 'test1', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
|
||||
assert result[1] == {'name': 'test2', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
|
||||
expected_pipeline = input_data['aggregation']
|
||||
expected_pipeline.append({'$project': {'_id': 0}})
|
||||
|
||||
mongo_db.mongo_db_repository.aggregate.assert_called_once_with(
|
||||
'test_collection', expected_pipeline, {}
|
||||
)
|
||||
|
||||
|
||||
def test_aggregate_documents_in_mongodb_failure(mongo_db):
|
||||
input_data = {
|
||||
'collection': 'test_collection',
|
||||
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
|
||||
}
|
||||
mongo_db.mongo_db_repository.aggregate = MagicMock(side_effect=Exception('Error'))
|
||||
|
||||
try:
|
||||
mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata})
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_AGGREGATION_ERROR',
|
||||
message='Failed to execute MongoDB aggregation: Error',
|
||||
block='aggregate_documents_in_mongodb',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_aggregate_documents_in_mongodb_missing_collection(mongo_db):
|
||||
input_data = {'query': {'aggregation': []}}
|
||||
|
||||
try:
|
||||
mongo_db.aggregate_documents_in_mongodb(input_data)
|
||||
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Collection name must be provided in the query.'
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected a ValueError to be raised')
|
||||
|
||||
|
||||
def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db):
|
||||
input_data = {'query': {'collection': 'test_collection'}}
|
||||
|
||||
try:
|
||||
mongo_db.aggregate_documents_in_mongodb(input_data)
|
||||
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Aggregation must be provided.'
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected a ValueError to be raised')
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_update_pipelines_timestamps_success(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'updated_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
]
|
||||
}
|
||||
mongo_db.mongo_db_repository.update_many = MagicMock(return_value=MagicMock())
|
||||
mongo_db.update_pipelines_timestamps(input_data)
|
||||
mongo_db.mongo_db_repository.update_many.assert_called_once_with(
|
||||
'orchestrated_schedules',
|
||||
{
|
||||
'$or': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1'},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2'},
|
||||
]
|
||||
},
|
||||
{'$set': {'updated_at': now_mock.return_value}},
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'updated_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
],
|
||||
**metadata,
|
||||
}
|
||||
|
||||
mongo_db.mongo_db_repository.update_many.side_effect = Exception('Error')
|
||||
try:
|
||||
mongo_db.update_pipelines_timestamps(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
|
||||
message='Failed to update pipelines timestamps: Error',
|
||||
block='update_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_create_pipelines_timestamps_success(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'created_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
]
|
||||
}
|
||||
mongo_db.mongo_db_repository.insert_many = MagicMock(return_value=MagicMock())
|
||||
mongo_db.create_pipelines_timestamps(input_data)
|
||||
mongo_db.mongo_db_repository.insert_many.assert_called_once_with(
|
||||
'orchestrated_schedules',
|
||||
[
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'updated_at': now_mock.return_value},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'updated_at': now_mock.return_value},
|
||||
],
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'created_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
],
|
||||
**metadata,
|
||||
}
|
||||
mongo_db.mongo_db_repository.insert_many.side_effect = Exception('Error')
|
||||
try:
|
||||
mongo_db.create_pipelines_timestamps(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_CREATE_PIPELINES_ERROR',
|
||||
message='Failed to create pipelines timestamps: Error',
|
||||
block='create_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'deleted_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
]
|
||||
}
|
||||
mongo_db.mongo_db_repository.delete_many = MagicMock(return_value=MagicMock())
|
||||
mongo_db.delete_pipelines_timestamps(input_data)
|
||||
mongo_db.mongo_db_repository.delete_many.assert_called_once_with(
|
||||
'orchestrated_schedules',
|
||||
{
|
||||
'$or': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1'},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2'},
|
||||
]
|
||||
},
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
@patch('orchestrator.activities.mongo_db.now')
|
||||
def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
|
||||
input_data = {
|
||||
'deleted_pipelines': [
|
||||
{'schedule_name': 'test1', 'namespace': 'test1', 'success': True},
|
||||
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
|
||||
],
|
||||
**metadata,
|
||||
}
|
||||
mongo_db.mongo_db_repository.delete_many.side_effect = Exception('Error')
|
||||
try:
|
||||
mongo_db.delete_pipelines_timestamps(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_DELETE_PIPELINES_ERROR',
|
||||
message='Failed to delete pipelines timestamps: Error',
|
||||
block='delete_pipelines_timestamps',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_create_collection_with_ttl_index_success(mongo_db):
|
||||
input_data = {
|
||||
**metadata,
|
||||
'pipelines': {
|
||||
'scouter-pipeline': {'topic': 'raw_scouter_pipeline'},
|
||||
'scouter-pipeline-2': {'topic': 'raw_scouter_pipeline_2'},
|
||||
'scouter-pipeline-3': {'topic': 'raw_scouter_pipeline_3'},
|
||||
},
|
||||
}
|
||||
|
||||
mongo_db.mongo_db_repository.database.list_collection_names = MagicMock(
|
||||
return_value=[
|
||||
'raw_scouter_pipeline_2',
|
||||
'raw_scouter_pipeline_3',
|
||||
]
|
||||
)
|
||||
|
||||
collection_1 = MagicMock(
|
||||
list_indexes=MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
'key': 'asdad',
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
collection_2 = MagicMock(
|
||||
list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': None}])
|
||||
)
|
||||
collection_3 = MagicMock(
|
||||
list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': 3600}])
|
||||
)
|
||||
|
||||
mongo_db.mongo_db_repository.database.__getitem__ = MagicMock(
|
||||
side_effect=[collection_1, collection_2, collection_3]
|
||||
)
|
||||
|
||||
mongo_db.create_collection_with_ttl_index(input_data)
|
||||
|
||||
mongo_db.mongo_db_repository.database.list_collection_names.assert_called_once_with()
|
||||
|
||||
mongo_db.mongo_db_repository.database.create_collection.assert_called_once_with(
|
||||
'raw_scouter_pipeline'
|
||||
)
|
||||
|
||||
collection_1.list_indexes.assert_called_once()
|
||||
collection_1.create_index.assert_called_once_with(
|
||||
'inserted_at', expireAfterSeconds=3600, background=True
|
||||
)
|
||||
|
||||
collection_2.list_indexes.assert_called_once()
|
||||
collection_2.create_index.assert_called_once_with(
|
||||
'inserted_at', expireAfterSeconds=3600, background=True
|
||||
)
|
||||
|
||||
collection_3.list_indexes.assert_called_once()
|
||||
collection_3.create_index.assert_not_called()
|
||||
|
||||
|
||||
def test_create_collection_with_ttl_index_failure(mongo_db):
|
||||
input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}}
|
||||
|
||||
mongo_db.mongo_db_repository.database.list_collection_names.return_value = []
|
||||
|
||||
mongo_db.mongo_db_repository.database.create_collection.side_effect = Exception('Error')
|
||||
|
||||
try:
|
||||
mongo_db.create_collection_with_ttl_index(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error'
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='MONGODB_CREATE_COLLECTION_ERROR',
|
||||
message='Failed to create collection raw_scouter_pipeline with TTL index: Error',
|
||||
block='create_collection_with_ttl_index',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_load_latest_data_none_last_data_timestamp(mongo_db):
|
||||
"""Test load_latest_data"""
|
||||
mongo_db.mongo_db_repository.find = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
'name': 'test1',
|
||||
'value': 1,
|
||||
'timestamp': '2023-01-01 12:00:00+0000',
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
result = mongo_db.load_latest_data(
|
||||
{
|
||||
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
'collection_name': 'test_collection',
|
||||
'last_data_timestamp': None,
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
}
|
||||
)
|
||||
|
||||
mongo_db.mongo_db_repository.find.assert_called_once_with(
|
||||
'test_collection',
|
||||
{'level': 'ERROR'},
|
||||
{'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
)
|
||||
|
||||
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00+0000'}]
|
||||
|
||||
|
||||
def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
|
||||
"""Test load_latest_data"""
|
||||
|
||||
mongo_db.mongo_db_repository.find = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
'name': 'test1',
|
||||
'value': 1,
|
||||
'timestamp': '2023-01-01 12:00:00+0000',
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
result = mongo_db.load_latest_data(
|
||||
{
|
||||
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
'collection_name': 'test_collection',
|
||||
'last_data_timestamp': '2023-01-01 12:00:00+0000',
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
}
|
||||
)
|
||||
|
||||
mongo_db.mongo_db_repository.find.assert_called_once_with(
|
||||
'test_collection',
|
||||
{
|
||||
'level': 'ERROR',
|
||||
'timestamp': {'$gt': '2023-01-01 12:00:00+0000'},
|
||||
},
|
||||
{'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
)
|
||||
|
||||
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00+0000'}]
|
||||
|
||||
|
||||
def test_load_latest_data_error(mongo_db):
|
||||
"""Test load_latest_data"""
|
||||
mongo_db.mongo_db_repository.find.side_effect = Exception('test')
|
||||
|
||||
try:
|
||||
mongo_db.load_latest_data(
|
||||
{
|
||||
'metadata': {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
'collection_name': 'test_collection',
|
||||
'last_data_timestamp': '2023-01-01 12:00:00+0000',
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
assert str(e) == 'test'
|
||||
|
||||
mongo_db.send_notification.assert_called_once_with(
|
||||
metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
|
||||
notification_id='MONGO_LOAD_ERROR',
|
||||
message='Error loading data from MongoDB: test',
|
||||
block='load_latest_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
409
tests/orchestrator/activities/test_slot_manager.py
Normal file
409
tests/orchestrator/activities/test_slot_manager.py
Normal file
@@ -0,0 +1,409 @@
|
||||
from datetime import timedelta
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pandas import DataFrame
|
||||
from pytest import fixture
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
|
||||
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test_schedule_name',
|
||||
'workflow_name': 'test_workflow_name',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('orchestrator.activities.slot_manager.RedisRepository')
|
||||
def slot_manager(_redis_mock):
|
||||
slot_manager = SlotManager(
|
||||
host='localhost',
|
||||
port=6379,
|
||||
username='admin',
|
||||
password='password',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
slot_manager.redis_repository = MagicMock()
|
||||
slot_manager.logger = MagicMock()
|
||||
slot_manager.notification_handler = MagicMock()
|
||||
slot_manager.send_notification = MagicMock()
|
||||
slot_manager.emit_metric = AsyncMock()
|
||||
|
||||
return slot_manager
|
||||
|
||||
|
||||
def test_load_opc_slots_no_slot_keys(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(return_value=[])
|
||||
assert slot_manager.load_opc_slots(metadata) == {}
|
||||
|
||||
|
||||
def test_load_opc_slots(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(
|
||||
return_value=[
|
||||
b'slot:opc_tags:1',
|
||||
b'slot:opc_tags:2',
|
||||
b'slot:opc_tags:3',
|
||||
]
|
||||
)
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(side_effect=['value1', 'value2', None])
|
||||
|
||||
response = slot_manager.load_opc_slots(metadata)
|
||||
|
||||
assert response == {
|
||||
'slot:opc_tags:1': 'value1',
|
||||
'slot:opc_tags:2': 'value2',
|
||||
'slot:opc_tags:3': None,
|
||||
}
|
||||
|
||||
|
||||
def test_load_opc_slots_no_decode(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(
|
||||
return_value=[
|
||||
'slot:opc_tags:1',
|
||||
'slot:opc_tags:2',
|
||||
'slot:opc_tags:3',
|
||||
]
|
||||
)
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(side_effect=['value1', 'value2', None])
|
||||
|
||||
response = slot_manager.load_opc_slots(metadata)
|
||||
|
||||
assert response == {
|
||||
'slot:opc_tags:1': 'value1',
|
||||
'slot:opc_tags:2': 'value2',
|
||||
'slot:opc_tags:3': None,
|
||||
}
|
||||
|
||||
|
||||
def test_load_opc_slots_error(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(
|
||||
return_value=[
|
||||
'slot:opc_tags:1',
|
||||
'slot:opc_tags:2',
|
||||
'slot:opc_tags:3',
|
||||
]
|
||||
)
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(side_effect=Exception('Test exception'))
|
||||
|
||||
try:
|
||||
slot_manager.load_opc_slots(metadata)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Test exception'
|
||||
slot_manager.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message='Failed to load OPC slots: Test exception',
|
||||
block='load_opc_slots',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_load_active_ingestors(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(
|
||||
return_value=[
|
||||
b'heartbeat:ingestor:1',
|
||||
b'heartbeat:ingestor:2',
|
||||
'heartbeat:ingestor:3',
|
||||
]
|
||||
)
|
||||
|
||||
response = slot_manager.load_active_ingestors(metadata)
|
||||
|
||||
assert response == ['heartbeat:ingestor:1', 'heartbeat:ingestor:2', 'heartbeat:ingestor:3']
|
||||
|
||||
|
||||
def test_load_active_ingestors_error(slot_manager):
|
||||
slot_manager.redis_repository.keys = MagicMock(
|
||||
return_value=[
|
||||
'heartbeat:ingestor:1',
|
||||
'heartbeat:ingestor:2',
|
||||
'heartbeat:ingestor:3',
|
||||
]
|
||||
)
|
||||
|
||||
slot_manager.redis_repository.keys = MagicMock(side_effect=Exception('Test exception'))
|
||||
|
||||
try:
|
||||
slot_manager.load_active_ingestors(metadata)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Test exception'
|
||||
slot_manager.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message='Failed to load active ingestors: Test exception',
|
||||
block='load_active_ingestors',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
def test_update_slots(slot_manager):
|
||||
slot_manager.redis_repository.set = MagicMock(side_effect=[None, Exception('Test exception')])
|
||||
|
||||
response = slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}})
|
||||
|
||||
slot_manager.redis_repository.set.assert_has_calls(
|
||||
[call('slot:opc_tags:1', 'value1', ttl=None), call('slot:opc_tags:2', 'value2', ttl=None)]
|
||||
)
|
||||
|
||||
assert response == {
|
||||
'1': {'success': True, 'message': 'Slot updated successfully'},
|
||||
'2': {'success': False, 'message': 'Test exception'},
|
||||
}
|
||||
|
||||
|
||||
def test_delete_slots(slot_manager):
|
||||
slot_manager.redis_repository.delete = MagicMock(
|
||||
side_effect=[None, Exception('Test exception')]
|
||||
)
|
||||
|
||||
response = slot_manager.delete_slots({'to_delete': ['1', '2']})
|
||||
|
||||
slot_manager.redis_repository.delete.assert_has_calls(
|
||||
[call('slot:opc_tags:1'), call('slot:opc_tags:2')]
|
||||
)
|
||||
|
||||
assert response == {
|
||||
'1': {'success': True, 'message': 'Slot deleted successfully'},
|
||||
'2': {'success': False, 'message': 'Test exception'},
|
||||
}
|
||||
|
||||
|
||||
def test_get_last_data_timestamp_none(slot_manager):
|
||||
"""Test get_last_data_timestamp"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(return_value=None)
|
||||
|
||||
result = slot_manager.get_last_data_timestamp(test_data)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_get_last_data_timestamp_not_none(slot_manager):
|
||||
"""Test get_last_data_timestamp"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(return_value='2023-01-01 12:00:00')
|
||||
|
||||
result = slot_manager.get_last_data_timestamp(test_data)
|
||||
|
||||
slot_manager.redis_repository.get.assert_called_once_with(
|
||||
'notification_last_timestamp:test_mail_type'
|
||||
)
|
||||
|
||||
assert result == '2023-01-01 12:00:00'
|
||||
|
||||
|
||||
def test_get_last_data_timestamp_error(slot_manager):
|
||||
"""Test get_last_data_timestamp error"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.get = MagicMock(side_effect=Exception('test'))
|
||||
|
||||
try:
|
||||
slot_manager.get_last_data_timestamp(test_data)
|
||||
|
||||
except Exception as e:
|
||||
assert str(e) == 'test'
|
||||
|
||||
slot_manager.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='REDIS_GET_ERROR',
|
||||
message='Error getting last data timestamp: test',
|
||||
block='get_last_data_timestamp',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected exception')
|
||||
|
||||
|
||||
def test_put_last_data_timestamp_empty_dataframe(slot_manager):
|
||||
"""Test put_last_data_timestamp with empty dataframe"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'data': DataFrame(columns=['name', 'value', 'timestamp']).to_dict('records'),
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.set = MagicMock()
|
||||
|
||||
result = slot_manager.put_last_data_timestamp(test_data)
|
||||
|
||||
assert result is None
|
||||
|
||||
slot_manager.set.assert_not_called()
|
||||
|
||||
|
||||
def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
|
||||
"""Test put_last_data_timestamp with not empty dataframe"""
|
||||
|
||||
data = DataFrame(
|
||||
{
|
||||
'name': ['sensor1', 'sensor2'],
|
||||
'value': [25.5, 30.0],
|
||||
'timestamp': ['2023-01-01 12:00:00', '2023-01-01 12:00:01'],
|
||||
}
|
||||
)
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'data': data.to_dict('records'),
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.set = MagicMock()
|
||||
|
||||
result = slot_manager.put_last_data_timestamp(test_data)
|
||||
|
||||
assert result == '2023-01-01 12:00:01'
|
||||
|
||||
slot_manager.redis_repository.set.assert_called_once_with(
|
||||
'notification_last_timestamp:test_mail_type', '2023-01-01 12:00:01', ttl=18000
|
||||
)
|
||||
|
||||
|
||||
def test_put_last_data_timestamp_error(slot_manager):
|
||||
"""Test put_last_data_timestamp error"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'workflow_name': 'test_pipeline',
|
||||
'schedule_name': 'test_schedule',
|
||||
'data': DataFrame(
|
||||
{
|
||||
'name': ['sensor1', 'sensor2'],
|
||||
'value': [25.5, 30.0],
|
||||
'timestamp': ['2023-01-01 12:00:00'] * 2,
|
||||
}
|
||||
).to_dict('records'),
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.set = MagicMock(side_effect=Exception('test'))
|
||||
|
||||
try:
|
||||
slot_manager.put_last_data_timestamp(test_data)
|
||||
|
||||
except Exception as e:
|
||||
assert str(e) == 'test'
|
||||
|
||||
slot_manager.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='REDIS_SET_ERROR',
|
||||
message='Error setting last data timestamp: test',
|
||||
block='put_last_data_timestamp',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected exception')
|
||||
|
||||
|
||||
def test_filter_notification_alerts(slot_manager):
|
||||
slot_manager.redis_repository.get = MagicMock(
|
||||
side_effect=[
|
||||
None,
|
||||
(now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||
now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||
None,
|
||||
(now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||
now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
||||
]
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'notification_package': [
|
||||
{'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'},
|
||||
{'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
|
||||
{'trigger': 'test_trigger_3', 'notification_id': 'test_notification_id_3'},
|
||||
],
|
||||
'sending_configs': [
|
||||
{'group_name': 'test_group_1', 'contents': ['core_alerts', 'persistent_alerts']},
|
||||
{'group_name': 'test_group_2', 'contents': ['core_alerts']},
|
||||
],
|
||||
'notification_ttl': 300,
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
response = slot_manager.filter_notification_alerts(input_data)
|
||||
|
||||
assert response == {
|
||||
'test_group_1': {
|
||||
'group_name': 'test_group_1',
|
||||
'contents': ['core_alerts', 'persistent_alerts'],
|
||||
'notifications': [
|
||||
{'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'},
|
||||
{'trigger': 'test_trigger_2', 'notification_id': 'test_notification_id_2'},
|
||||
],
|
||||
},
|
||||
'test_group_2': {
|
||||
'group_name': 'test_group_2',
|
||||
'contents': ['core_alerts'],
|
||||
'notifications': [
|
||||
{'trigger': 'test_trigger_1', 'notification_id': 'test_notification_id_1'}
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_store_notification_cache(slot_manager):
|
||||
"""Test store_notification_cache"""
|
||||
test_data = {
|
||||
**metadata,
|
||||
'log_report': DataFrame(
|
||||
{
|
||||
'status': ['sent', 'error'],
|
||||
'schedule': ['test_schedule_1', 'test_schedule_2'],
|
||||
'notification_id': ['test_notification_id_1', 'test_notification_id_2'],
|
||||
}
|
||||
).to_dict(),
|
||||
'sent_ttl': 600,
|
||||
}
|
||||
|
||||
slot_manager.redis_repository.set = MagicMock()
|
||||
|
||||
slot_manager.store_notification_cache(test_data)
|
||||
|
||||
slot_manager.redis_repository.set.assert_called_once_with(
|
||||
'test_schedule_1:test_notification_id_1', ANY, ttl=600
|
||||
)
|
||||
686
tests/orchestrator/activities/test_temporal_manager.py
Normal file
686
tests/orchestrator/activities/test_temporal_manager.py
Normal file
@@ -0,0 +1,686 @@
|
||||
from datetime import timedelta
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test_schedule_name',
|
||||
'workflow_name': 'test_workflow_name',
|
||||
'model_name': 'test_model_name',
|
||||
'model_id': 'test_model_id',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@fixture
|
||||
def temporal_manager():
|
||||
temporal_manager = TemporalManager(
|
||||
host='localhost:7233',
|
||||
scouter_namespace='scouter',
|
||||
laborious_namespace='laborious',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'] = MagicMock()
|
||||
temporal_manager.temporal_clients['laborious'] = MagicMock()
|
||||
temporal_manager.send_notification_async = AsyncMock()
|
||||
temporal_manager.emit_metric = AsyncMock()
|
||||
temporal_manager.send_notification = MagicMock()
|
||||
|
||||
return temporal_manager
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.Client.connect', new_callable=AsyncMock)
|
||||
async def test_connect_to_temporal(connect_mock, temporal_manager):
|
||||
await temporal_manager.connect_to_temporal()
|
||||
connect_mock.assert_has_calls(
|
||||
[
|
||||
call(target_host='localhost:7233', namespace='scouter'),
|
||||
call(target_host='localhost:7233', namespace='laborious'),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def async_iter():
|
||||
yield MagicMock(id='test-schedule-id', search_attributes={'orchestrated': ['true']})
|
||||
yield MagicMock(id='test-schedule-id-2', search_attributes={'Attr': ['false']})
|
||||
yield MagicMock(id='test-schedule-id-3', search_attributes={'Attr': ['false']})
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_normalize_schedules(temporal_manager):
|
||||
input_data = {
|
||||
'orchestrated_schedules': {
|
||||
'scouter': {'test-scouter': '2021-01-01'},
|
||||
'laborious': {'test-schedule-id1': '2021-01-01'},
|
||||
},
|
||||
**metadata,
|
||||
}
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
|
||||
return_value=async_iter()
|
||||
)
|
||||
temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock(
|
||||
return_value=async_iter()
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
|
||||
return_value=MagicMock(delete=AsyncMock())
|
||||
)
|
||||
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
|
||||
return_value=MagicMock(delete=AsyncMock())
|
||||
)
|
||||
|
||||
await temporal_manager.normalize_schedules(input_data)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls(
|
||||
[
|
||||
call('test-schedule-id'),
|
||||
]
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients[
|
||||
'scouter'
|
||||
].get_schedule_handle.return_value.delete.assert_awaited_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_normalize_schedules_error(temporal_manager):
|
||||
temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock(
|
||||
side_effect=Exception('Test exception')
|
||||
)
|
||||
|
||||
try:
|
||||
await temporal_manager.normalize_schedules(metadata)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Test exception'
|
||||
temporal_manager.send_notification_async.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR',
|
||||
message='Failed to normalize schedules: Test exception',
|
||||
block='normalize_schedules',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY,
|
||||
)
|
||||
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
|
||||
@patch('orchestrator.activities.temporal_manager.Schedule')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
|
||||
@patch('orchestrator.activities.temporal_manager.SearchAttributePair')
|
||||
async def test_create_schedule(
|
||||
mock_search_attribute_pair,
|
||||
mock_typed_search_attributes,
|
||||
mock_schedule_spec,
|
||||
mock_schedule_interval_spec,
|
||||
mock_schedule_action_start_workflow,
|
||||
mock_schedule,
|
||||
mock_parse_frequency,
|
||||
temporal_manager,
|
||||
):
|
||||
input_data = {
|
||||
'schedules': {
|
||||
'scouter': {
|
||||
'test-schedule': {
|
||||
'model_id': 1,
|
||||
'model_name': 'test-model-name',
|
||||
'workflow_type': 'test-workflow',
|
||||
'frequency': '1m',
|
||||
'data': {'test': 'test'},
|
||||
'execution_timeout_seconds': 100,
|
||||
'task_timeout_seconds': 100,
|
||||
},
|
||||
'test-schedule-invalid-frequency': {
|
||||
'model_id': 2,
|
||||
'model_name': 'test-model-name',
|
||||
'workflow_type': 'test-workflow',
|
||||
'frequency': '10y',
|
||||
'offset': '5m',
|
||||
'data': {'test': 'test'},
|
||||
'execution_timeout_seconds': 400,
|
||||
'task_timeout_seconds': 400,
|
||||
},
|
||||
},
|
||||
'laborious': {
|
||||
'test-schedule-laborious': {
|
||||
'model_id': 1,
|
||||
'model_name': 'test-model-name',
|
||||
'workflow_type': 'test-workflow',
|
||||
'frequency': '2m',
|
||||
'offset': '1h',
|
||||
'data': {'test': 'test'},
|
||||
'execution_timeout_seconds': 500,
|
||||
'task_timeout_seconds': 500,
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
|
||||
temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock()
|
||||
|
||||
report = await temporal_manager.create_schedules(input_data)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with(
|
||||
'test-schedule',
|
||||
mock_schedule.return_value,
|
||||
search_attributes=mock_typed_search_attributes.return_value,
|
||||
)
|
||||
temporal_manager.temporal_clients['laborious'].create_schedule.assert_called_once_with(
|
||||
'test-schedule-laborious',
|
||||
mock_schedule.return_value,
|
||||
search_attributes=mock_typed_search_attributes.return_value,
|
||||
)
|
||||
|
||||
mock_schedule.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
action=mock_schedule_action_start_workflow.return_value,
|
||||
spec=mock_schedule_spec.return_value,
|
||||
),
|
||||
call(
|
||||
action=mock_schedule_action_start_workflow.return_value,
|
||||
spec=mock_schedule_spec.return_value,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_schedule_action_start_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'test-workflow',
|
||||
input_data['schedules']['scouter']['test-schedule'],
|
||||
id='test-schedule',
|
||||
task_queue='test-workflow-queue',
|
||||
execution_timeout=timedelta(seconds=100),
|
||||
run_timeout=timedelta(seconds=100),
|
||||
task_timeout=timedelta(seconds=100),
|
||||
typed_search_attributes=mock_typed_search_attributes.return_value,
|
||||
),
|
||||
call(
|
||||
'test-workflow',
|
||||
input_data['schedules']['scouter']['test-schedule-invalid-frequency'],
|
||||
id='test-schedule-invalid-frequency',
|
||||
task_queue='test-workflow-queue',
|
||||
execution_timeout=timedelta(seconds=400),
|
||||
run_timeout=timedelta(seconds=400),
|
||||
task_timeout=timedelta(seconds=400),
|
||||
typed_search_attributes=mock_typed_search_attributes.return_value,
|
||||
),
|
||||
call(
|
||||
'test-workflow',
|
||||
input_data['schedules']['laborious']['test-schedule-laborious'],
|
||||
id='test-schedule-laborious',
|
||||
task_queue='test-workflow-queue',
|
||||
execution_timeout=timedelta(seconds=500),
|
||||
run_timeout=timedelta(seconds=500),
|
||||
task_timeout=timedelta(seconds=500),
|
||||
typed_search_attributes=mock_typed_search_attributes.return_value,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_schedule_spec.assert_has_calls(
|
||||
[
|
||||
call(intervals=[mock_schedule_interval_spec.return_value]),
|
||||
call(intervals=[mock_schedule_interval_spec.return_value]),
|
||||
]
|
||||
)
|
||||
|
||||
mock_schedule_interval_spec.assert_has_calls(
|
||||
[
|
||||
call(every=timedelta(seconds=60), offset=timedelta(seconds=0)),
|
||||
call(every=timedelta(seconds=120), offset=timedelta(seconds=3600)),
|
||||
]
|
||||
)
|
||||
|
||||
mock_parse_frequency.assert_has_calls(
|
||||
[call('1m'), call('0m'), call('10y'), call('2m'), call('1h')]
|
||||
)
|
||||
|
||||
mock_typed_search_attributes.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
[
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
]
|
||||
),
|
||||
call(
|
||||
[
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
]
|
||||
),
|
||||
call(
|
||||
[
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
mock_search_attribute_pair.return_value,
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
mock_search_attribute_pair.assert_has_calls(
|
||||
[
|
||||
call(key=temporal_manager.model_id_id_key, value=1),
|
||||
call(key=temporal_manager.model_name_id_key, value='test-model-name'),
|
||||
call(key=temporal_manager.orchestrated_id_key, value='true'),
|
||||
call(key=temporal_manager.model_id_id_key, value=2),
|
||||
call(key=temporal_manager.model_name_id_key, value='test-model-name'),
|
||||
call(key=temporal_manager.orchestrated_id_key, value='true'),
|
||||
]
|
||||
)
|
||||
|
||||
assert report == [
|
||||
{
|
||||
'schedule_name': 'test-schedule',
|
||||
'namespace': 'scouter',
|
||||
'success': True,
|
||||
'message': 'Schedule created successfully',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule-invalid-frequency',
|
||||
'namespace': 'scouter',
|
||||
'success': False,
|
||||
'message': 'Invalid frequency',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule-laborious',
|
||||
'namespace': 'laborious',
|
||||
'success': True,
|
||||
'message': 'Schedule created successfully',
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_create_schedules_with_no_client(temporal_manager):
|
||||
temporal_manager.temporal_clients = {}
|
||||
input_data = {
|
||||
'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}}
|
||||
}
|
||||
|
||||
try:
|
||||
await temporal_manager.create_schedules(input_data)
|
||||
except Exception as e:
|
||||
assert (
|
||||
str(e)
|
||||
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
|
||||
async def test_update_schedules(
|
||||
_mock_schedule_interval_spec, _mock_parse_frequency, temporal_manager
|
||||
):
|
||||
input_mock = MagicMock(args=MagicMock())
|
||||
temporal_manager.schedule_handles = {
|
||||
'scouter': {
|
||||
'test-schedule': MagicMock(
|
||||
update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
|
||||
)
|
||||
},
|
||||
'laborious': {
|
||||
'test-schedule-laborious': MagicMock(
|
||||
update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
|
||||
)
|
||||
},
|
||||
}
|
||||
input_data = {
|
||||
'schedules': {
|
||||
'scouter': {
|
||||
'test-schedule': {
|
||||
'workflow_type': 'scouter',
|
||||
'frequency': '1m',
|
||||
'data': {'test': 'test'},
|
||||
},
|
||||
'test-schedule_no_handler': {
|
||||
'workflow_type': 'scouter',
|
||||
'frequency': '1m',
|
||||
'data': {'test': 'test'},
|
||||
},
|
||||
},
|
||||
'laborious': {
|
||||
'test-schedule-laborious': {
|
||||
'workflow_type': 'laborious',
|
||||
'frequency': '2m',
|
||||
'data': {'test': 'test'},
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
handler_scouter = MagicMock(
|
||||
update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
|
||||
)
|
||||
|
||||
handler_laborious = MagicMock(
|
||||
update=AsyncMock(update=AsyncMock(side_effect=lambda f: f(input_mock)))
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
|
||||
side_effect=[handler_scouter, None]
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
|
||||
side_effect=[handler_laborious]
|
||||
)
|
||||
|
||||
report = await temporal_manager.update_schedules(input_data)
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].get_schedule_handle.assert_has_calls(
|
||||
[call('test-schedule'), call('test-schedule_no_handler')]
|
||||
)
|
||||
temporal_manager.temporal_clients['laborious'].get_schedule_handle.assert_has_calls(
|
||||
[call('test-schedule-laborious')]
|
||||
)
|
||||
|
||||
handler_scouter.update.assert_called_once()
|
||||
handler_laborious.update.assert_called_once()
|
||||
|
||||
assert report == [
|
||||
{
|
||||
'schedule_name': 'test-schedule',
|
||||
'namespace': 'scouter',
|
||||
'success': True,
|
||||
'message': 'Schedule updated successfully',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule_no_handler',
|
||||
'namespace': 'scouter',
|
||||
'success': False,
|
||||
'message': 'Schedule test-schedule_no_handler not found',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule-laborious',
|
||||
'namespace': 'laborious',
|
||||
'success': True,
|
||||
'message': 'Schedule updated successfully',
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_schedules_with_no_client(temporal_manager):
|
||||
temporal_manager.temporal_clients = {}
|
||||
input_data = {
|
||||
'schedules': {'abc': {'test-schedule': {'frequency': '1m', 'data': {'test': 'test'}}}}
|
||||
}
|
||||
|
||||
try:
|
||||
await temporal_manager.update_schedules(input_data)
|
||||
except Exception as e:
|
||||
assert (
|
||||
str(e)
|
||||
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_schedules(temporal_manager):
|
||||
temporal_manager.schedule_handles = {
|
||||
'scouter': {'test-schedule': MagicMock(delete=AsyncMock())},
|
||||
'laborious': {'test-schedule-laborious': MagicMock(delete=AsyncMock())},
|
||||
}
|
||||
|
||||
input_data = {
|
||||
'schedules': {
|
||||
'scouter': ['test-schedule', 'test-schedule_no_handler'],
|
||||
'laborious': ['test-schedule-laborious'],
|
||||
}
|
||||
}
|
||||
handler_scouter = MagicMock(delete=AsyncMock())
|
||||
|
||||
handler_laborious = MagicMock(delete=AsyncMock())
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock(
|
||||
side_effect=[handler_scouter, None]
|
||||
)
|
||||
|
||||
temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock(
|
||||
side_effect=[handler_laborious]
|
||||
)
|
||||
|
||||
report = await temporal_manager.delete_schedules(input_data)
|
||||
|
||||
handler_scouter.delete.assert_called_once()
|
||||
handler_laborious.delete.assert_called_once()
|
||||
|
||||
assert report == [
|
||||
{
|
||||
'schedule_name': 'test-schedule',
|
||||
'namespace': 'scouter',
|
||||
'success': True,
|
||||
'message': 'Schedule deleted successfully',
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule_no_handler',
|
||||
'namespace': 'scouter',
|
||||
'success': False,
|
||||
'message': 'Schedule test-schedule_no_handler not found',
|
||||
'attachment': ANY,
|
||||
},
|
||||
{
|
||||
'schedule_name': 'test-schedule-laborious',
|
||||
'namespace': 'laborious',
|
||||
'success': True,
|
||||
'message': 'Schedule deleted successfully',
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_delete_schedules_with_no_client(temporal_manager):
|
||||
temporal_manager.temporal_clients = {}
|
||||
input_data = {'schedules': {'abc': ['test-schedule']}}
|
||||
|
||||
try:
|
||||
await temporal_manager.delete_schedules(input_data)
|
||||
except Exception as e:
|
||||
assert (
|
||||
str(e)
|
||||
== f'Temporal client for abc not found, clients: {temporal_manager.temporal_clients}'
|
||||
)
|
||||
|
||||
|
||||
def test_build_task_queue_name_non_runtime_workflow(temporal_manager):
|
||||
task_queue_name = temporal_manager._build_task_queue_name(
|
||||
'test-workflow', {'workflow_type': 'test-workflow'}
|
||||
)
|
||||
|
||||
assert task_queue_name == 'test-workflow-queue'
|
||||
|
||||
|
||||
def test_build_task_queue_name_default_runtime_legacy(temporal_manager):
|
||||
task_queue_name = temporal_manager._build_task_queue_name('drift', {'workflow_type': 'drift'})
|
||||
|
||||
assert task_queue_name == 'drift-legacy-queue'
|
||||
|
||||
|
||||
def test_build_task_queue_name_explicit_runtime(temporal_manager):
|
||||
task_queue_name = temporal_manager._build_task_queue_name(
|
||||
'drift', {'workflow_type': 'drift', 'runtime': 'tenant-x'}
|
||||
)
|
||||
|
||||
assert task_queue_name == 'drift-tenant-x-queue'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
|
||||
async def test_make_schedule_updater_patches_args_and_intervals(
|
||||
_mock_schedule_interval_spec, _mock_parse_frequency, temporal_manager
|
||||
):
|
||||
schedule = {'workflow_type': 'scouter', 'frequency': '2m', 'offset': '1h', 'data': 'new'}
|
||||
update_schedule = temporal_manager._make_schedule_updater(schedule, metadata)
|
||||
|
||||
schedule_action = MagicMock()
|
||||
input_data = MagicMock()
|
||||
input_data.description.schedule.action = schedule_action
|
||||
|
||||
result = await update_schedule(input_data)
|
||||
|
||||
assert schedule_action.args == [schedule]
|
||||
assert input_data.description.schedule.spec.intervals == [
|
||||
_mock_schedule_interval_spec.return_value
|
||||
]
|
||||
_mock_schedule_interval_spec.assert_called_once_with(
|
||||
every=timedelta(seconds=120), offset=timedelta(seconds=3600)
|
||||
)
|
||||
assert result.schedule == input_data.description.schedule
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_make_schedule_updater_skips_args_without_attribute(temporal_manager):
|
||||
schedule = {'workflow_type': 'scouter', 'frequency': '1m', 'data': 'new'}
|
||||
update_schedule = temporal_manager._make_schedule_updater(schedule, metadata)
|
||||
|
||||
input_data = MagicMock()
|
||||
input_data.description.schedule.action = object()
|
||||
|
||||
await update_schedule(input_data)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_single_schedule_not_found(temporal_manager):
|
||||
client = MagicMock(get_schedule_handle=MagicMock(return_value=None))
|
||||
|
||||
try:
|
||||
await temporal_manager._update_single_schedule(
|
||||
client, 'missing-schedule', {'workflow_type': 'scouter'}, metadata
|
||||
)
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Schedule missing-schedule not found'
|
||||
else:
|
||||
raise AssertionError('Expected a ValueError to be raised')
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_single_schedule_updates_handle(temporal_manager):
|
||||
handler = MagicMock(update=AsyncMock())
|
||||
client = MagicMock(get_schedule_handle=MagicMock(return_value=handler))
|
||||
schedule = {'workflow_type': 'drift', 'frequency': '1m'}
|
||||
|
||||
await temporal_manager._update_single_schedule(client, 'test-schedule', schedule, metadata)
|
||||
|
||||
client.get_schedule_handle.assert_called_once_with('test-schedule')
|
||||
handler.update.assert_awaited_once()
|
||||
assert schedule['task_queue'] == 'drift-legacy-queue'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
|
||||
@patch('orchestrator.activities.temporal_manager.Schedule')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
|
||||
@patch('orchestrator.activities.temporal_manager.SearchAttributePair')
|
||||
async def test_create_schedules_default_runtime_legacy_queue(
|
||||
_mock_search_attribute_pair,
|
||||
_mock_typed_search_attributes,
|
||||
_mock_schedule_spec,
|
||||
_mock_schedule_interval_spec,
|
||||
mock_schedule_action_start_workflow,
|
||||
_mock_schedule,
|
||||
_mock_parse_frequency,
|
||||
temporal_manager,
|
||||
):
|
||||
input_data = {
|
||||
'schedules': {
|
||||
'scouter': {
|
||||
'test-schedule': {
|
||||
'model_id': 1,
|
||||
'model_name': 'test-model-name',
|
||||
'workflow_type': 'drift',
|
||||
'frequency': '1m',
|
||||
'data': {'test': 'test'},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
|
||||
|
||||
await temporal_manager.create_schedules(input_data)
|
||||
|
||||
mock_schedule_action_start_workflow.assert_called_once_with(
|
||||
'drift',
|
||||
input_data['schedules']['scouter']['test-schedule'],
|
||||
id='test-schedule',
|
||||
task_queue='drift-legacy-queue',
|
||||
execution_timeout=timedelta(seconds=300),
|
||||
run_timeout=timedelta(seconds=300),
|
||||
task_timeout=timedelta(seconds=300),
|
||||
typed_search_attributes=_mock_typed_search_attributes.return_value,
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.activities.temporal_manager.parse_frequency', side_effect=parse_frequency)
|
||||
@patch('orchestrator.activities.temporal_manager.Schedule')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleActionStartWorkflow')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleIntervalSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.ScheduleSpec')
|
||||
@patch('orchestrator.activities.temporal_manager.TypedSearchAttributes')
|
||||
@patch('orchestrator.activities.temporal_manager.SearchAttributePair')
|
||||
async def test_create_schedules_tenant_runtime_queue(
|
||||
_mock_search_attribute_pair,
|
||||
_mock_typed_search_attributes,
|
||||
_mock_schedule_spec,
|
||||
_mock_schedule_interval_spec,
|
||||
mock_schedule_action_start_workflow,
|
||||
_mock_schedule,
|
||||
_mock_parse_frequency,
|
||||
temporal_manager,
|
||||
):
|
||||
input_data = {
|
||||
'schedules': {
|
||||
'scouter': {
|
||||
'test-schedule': {
|
||||
'model_id': 1,
|
||||
'model_name': 'test-model-name',
|
||||
'workflow_type': 'drift',
|
||||
'frequency': '1m',
|
||||
'runtime': 'tenant-x',
|
||||
'data': {'test': 'test'},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock()
|
||||
|
||||
await temporal_manager.create_schedules(input_data)
|
||||
|
||||
mock_schedule_action_start_workflow.assert_called_once_with(
|
||||
'drift',
|
||||
input_data['schedules']['scouter']['test-schedule'],
|
||||
id='test-schedule',
|
||||
task_queue='drift-tenant-x-queue',
|
||||
execution_timeout=timedelta(seconds=300),
|
||||
run_timeout=timedelta(seconds=300),
|
||||
task_timeout=timedelta(seconds=300),
|
||||
typed_search_attributes=_mock_typed_search_attributes.return_value,
|
||||
)
|
||||
0
tests/orchestrator/utils/__init__.py
Normal file
0
tests/orchestrator/utils/__init__.py
Normal file
174
tests/orchestrator/utils/test_connectors_config.py
Normal file
174
tests/orchestrator/utils/test_connectors_config.py
Normal file
@@ -0,0 +1,174 @@
|
||||
from os import environ
|
||||
|
||||
from orchestrator.utils.connectors_config import (
|
||||
build_couchbase_config,
|
||||
build_email_config,
|
||||
build_mongodb_config,
|
||||
build_postgres_config,
|
||||
build_redis_config,
|
||||
build_temporal_config,
|
||||
)
|
||||
|
||||
|
||||
def test_build_redis_config_with_env_vars():
|
||||
environ['REDIS_HOST'] = 'localhost'
|
||||
environ['REDIS_PORT'] = '6379'
|
||||
environ['REDIS_USERNAME'] = 'sientia'
|
||||
environ['REDIS_PASSWORD'] = 'sientia'
|
||||
assert build_redis_config() == {
|
||||
'host': 'localhost',
|
||||
'port': 6379,
|
||||
'username': 'sientia',
|
||||
'password': 'sientia',
|
||||
}
|
||||
|
||||
|
||||
def test_build_couchbase_config_with_env_vars():
|
||||
environ['COUCHBASE_CONNECTION_STRING'] = 'couchbase://localhost'
|
||||
environ['COUCHBASE_USERNAME'] = 'sientia'
|
||||
environ['COUCHBASE_PASSWORD'] = 'sientia'
|
||||
assert build_couchbase_config() == {
|
||||
'connection_string': 'couchbase://localhost',
|
||||
'username': 'sientia',
|
||||
'password': 'sientia',
|
||||
}
|
||||
|
||||
|
||||
def test_build_redis_config_with_defaults():
|
||||
environ.pop('REDIS_HOST', None)
|
||||
environ.pop('REDIS_PORT', None)
|
||||
environ.pop('REDIS_USERNAME', None)
|
||||
environ.pop('REDIS_PASSWORD', None)
|
||||
assert build_redis_config() == {
|
||||
'host': 'localhost',
|
||||
'port': 6379,
|
||||
'username': 'default',
|
||||
'password': 'bdnZOpcyiL',
|
||||
}
|
||||
|
||||
|
||||
def test_build_couchbase_config_with_defaults():
|
||||
environ.pop('COUCHBASE_CONNECTION_STRING', None)
|
||||
environ.pop('COUCHBASE_USERNAME', None)
|
||||
environ.pop('COUCHBASE_PASSWORD', None)
|
||||
assert build_couchbase_config() == {
|
||||
'connection_string': 'couchbase://localhost',
|
||||
'username': 'sientia',
|
||||
'password': 'sientia',
|
||||
}
|
||||
|
||||
|
||||
def test_build_mongo_db_config_with_env_vars():
|
||||
environ['MONGODB_USERNAME'] = 'sientia1'
|
||||
environ['MONGODB_PASSWORD'] = 'sientia1'
|
||||
environ['MONGODB_URL'] = 'localhost:27018'
|
||||
environ['MONGODB_DATABASE_NAME'] = 'test_db'
|
||||
environ['MONGODB_TTL_INDEX_HOURS'] = '2'
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
|
||||
'database_name': 'test_db',
|
||||
'ttl_index_seconds': 7200,
|
||||
}
|
||||
|
||||
|
||||
def test_build_mongo_db_config_with_defaults():
|
||||
environ.pop('MONGODB_USERNAME', None)
|
||||
environ.pop('MONGODB_PASSWORD', None)
|
||||
environ.pop('MONGODB_DATABASE_NAME', None)
|
||||
environ.pop('MONGODB_URL', None)
|
||||
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
|
||||
'database_name': 'sientia',
|
||||
'ttl_index_seconds': 3600,
|
||||
}
|
||||
|
||||
|
||||
def test_build_temporal_config_with_env_vars():
|
||||
environ['TEMPORAL_HOST'] = 'localhost:7233'
|
||||
environ['TEMPORAL_SCOUTER_NAMESPACE'] = 'scouter'
|
||||
environ['TEMPORAL_LABORIOUS_NAMESPACE'] = 'laborious'
|
||||
assert build_temporal_config() == {
|
||||
'temporal_host': 'localhost:7233',
|
||||
'temporal_namespace': 'default',
|
||||
'temporal_scouter_namespace': 'scouter',
|
||||
'temporal_laborious_namespace': 'laborious',
|
||||
}
|
||||
|
||||
|
||||
def test_build_temporal_config_with_defaults():
|
||||
environ.pop('TEMPORAL_HOST', None)
|
||||
environ.pop('TEMPORAL_SCOUTER_NAMESPACE', None)
|
||||
environ.pop('TEMPORAL_LABORIOUS_NAMESPACE', None)
|
||||
assert build_temporal_config() == {
|
||||
'temporal_host': 'localhost:7233',
|
||||
'temporal_namespace': 'default',
|
||||
'temporal_scouter_namespace': 'scouter',
|
||||
'temporal_laborious_namespace': 'laborious',
|
||||
}
|
||||
|
||||
|
||||
def test_build_email_config_with_env_vars():
|
||||
environ['EMAIL_SENDER'] = 'test@test.com'
|
||||
environ['EMAIL_SENDER_PASSWORD'] = 'test'
|
||||
environ['EMAIL_SMTP_SERVER'] = 'test'
|
||||
environ['EMAIL_SMTP_PORT'] = '587'
|
||||
assert build_email_config() == {
|
||||
'sender_email': 'test@test.com',
|
||||
'sender_password': 'test',
|
||||
'smtp_server': 'test',
|
||||
'smtp_port': 587,
|
||||
}
|
||||
|
||||
|
||||
def test_build_email_config_with_defaults():
|
||||
environ.pop('EMAIL_SENDER', None)
|
||||
environ.pop('EMAIL_SENDER_PASSWORD', None)
|
||||
environ.pop('EMAIL_SMTP_SERVER', None)
|
||||
environ.pop('EMAIL_SMTP_PORT', None)
|
||||
|
||||
assert build_email_config() == {
|
||||
'sender_email': 'sientia-alerts@aignosi.com',
|
||||
'sender_password': 'sientia',
|
||||
'smtp_server': None,
|
||||
'smtp_port': 587,
|
||||
}
|
||||
|
||||
|
||||
def test_build_postgres_config_with_env_vars():
|
||||
environ['POSTGRES_HOST'] = 'localhost'
|
||||
environ['POSTGRES_PORT'] = '5432'
|
||||
environ['POSTGRES_USER'] = 'sientia'
|
||||
environ['POSTGRES_PASSWORD'] = 'sientia'
|
||||
environ['POSTGRES_DBNAME'] = 'sientia'
|
||||
environ['POSTGRES_MIN_CONNECTIONS'] = '5'
|
||||
environ['POSTGRES_MAX_CONNECTIONS'] = '20'
|
||||
assert build_postgres_config() == {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'sientia',
|
||||
'password': 'sientia',
|
||||
'dbname': 'sientia',
|
||||
'min_connections': 5,
|
||||
'max_connections': 20,
|
||||
}
|
||||
|
||||
|
||||
def test_build_postgres_config_with_defaults():
|
||||
environ.pop('POSTGRES_HOST', None)
|
||||
environ.pop('POSTGRES_PORT', None)
|
||||
environ.pop('POSTGRES_USER', None)
|
||||
environ.pop('POSTGRES_PASSWORD', None)
|
||||
environ.pop('POSTGRES_DBNAME', None)
|
||||
environ.pop('POSTGRES_MIN_CONNECTIONS', None)
|
||||
environ.pop('POSTGRES_MAX_CONNECTIONS', None)
|
||||
|
||||
assert build_postgres_config() == {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'sientia',
|
||||
'password': 'sientia',
|
||||
'dbname': 'sientia',
|
||||
'min_connections': 5,
|
||||
'max_connections': 20,
|
||||
}
|
||||
15
tests/orchestrator/utils/test_converters.py
Normal file
15
tests/orchestrator/utils/test_converters.py
Normal file
@@ -0,0 +1,15 @@
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
|
||||
def test_parse_frequency():
|
||||
assert parse_frequency('1s') == 1
|
||||
assert parse_frequency('1m') == 60
|
||||
assert parse_frequency('1h') == 60 * 60
|
||||
assert parse_frequency('1d') == 60 * 60 * 24
|
||||
|
||||
try:
|
||||
parse_frequency('1')
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Invalid frequency'
|
||||
else:
|
||||
raise AssertionError('Expected an exception to be raised')
|
||||
198
tests/orchestrator/utils/test_email_builder.py
Normal file
198
tests/orchestrator/utils/test_email_builder.py
Normal file
@@ -0,0 +1,198 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
from orchestrator.utils.email_builder import EmailBuilder
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('orchestrator.utils.email_builder.open')
|
||||
def report_builder(open_mock):
|
||||
return EmailBuilder(MagicMock())
|
||||
|
||||
|
||||
@patch('orchestrator.utils.email_builder.Template')
|
||||
def test_replace_parameters(template, report_builder):
|
||||
output = report_builder.replace_parameters('template', {'key': 'value'})
|
||||
|
||||
assert output == template.return_value.render.return_value
|
||||
|
||||
template.assert_called_once_with('template')
|
||||
template.return_value.render.assert_called_once_with({'key': 'value'})
|
||||
|
||||
|
||||
def test_parameters(report_builder):
|
||||
report_builder.replace_parameters = MagicMock()
|
||||
general_events = {
|
||||
'ERROR': {
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}],
|
||||
},
|
||||
]
|
||||
},
|
||||
'WARNING': {
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [
|
||||
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}
|
||||
],
|
||||
},
|
||||
]
|
||||
},
|
||||
'INFO': {
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}],
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
output = report_builder.parameters(general_events, 'model_name')
|
||||
|
||||
assert output == {
|
||||
'mail_type': 'model_name',
|
||||
'error_events': report_builder.replace_parameters.return_value,
|
||||
'warning_events': report_builder.replace_parameters.return_value,
|
||||
'info_events': report_builder.replace_parameters.return_value,
|
||||
}
|
||||
|
||||
report_builder.replace_parameters.assert_any_call(
|
||||
report_builder.general_template,
|
||||
{
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [{'notification_id': 'ID_1', 'level': 'ERROR', 'project': 'project'}],
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
report_builder.replace_parameters.assert_any_call(
|
||||
report_builder.general_template,
|
||||
{
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [
|
||||
{'notification_id': 'ID_2', 'level': 'WARNING', 'project': 'project'}
|
||||
],
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
report_builder.replace_parameters.assert_any_call(
|
||||
report_builder.general_template,
|
||||
{
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [{'notification_id': 'ID_3', 'level': 'INFO', 'project': 'project'}],
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_build_email(report_builder):
|
||||
report_builder.parameters = MagicMock()
|
||||
report_builder.replace_parameters = MagicMock()
|
||||
|
||||
report_data = [
|
||||
{
|
||||
'notification_id': 'ID_1',
|
||||
'level': 'ERROR',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
{
|
||||
'notification_id': 'ID_2',
|
||||
'level': 'WARNING',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
{
|
||||
'notification_id': 'ID_2',
|
||||
'level': 'INFO',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
{
|
||||
'notification_id': 'ID_3',
|
||||
'level': 'ERROR',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
]
|
||||
|
||||
html = report_builder.build_email(report_data, 'type_1')
|
||||
|
||||
report_builder.replace_parameters.assert_called_once_with(
|
||||
report_builder.report_template, report_builder.parameters.return_value
|
||||
)
|
||||
|
||||
assert html == report_builder.replace_parameters.return_value
|
||||
|
||||
report_builder.parameters.assert_called_once_with(
|
||||
{
|
||||
'ERROR': {
|
||||
'section_name': 'Errors detected:',
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [
|
||||
{
|
||||
'notification_id': 'ID_1',
|
||||
'level': 'ERROR',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
{
|
||||
'notification_id': 'ID_3',
|
||||
'level': 'ERROR',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
'WARNING': {
|
||||
'section_name': 'Warnings detected:',
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [
|
||||
{
|
||||
'notification_id': 'ID_2',
|
||||
'level': 'WARNING',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
'INFO': {
|
||||
'section_name': 'Infos detected:',
|
||||
'models': [
|
||||
{
|
||||
'model_name': 'model_name',
|
||||
'events': [
|
||||
{
|
||||
'notification_id': 'ID_2',
|
||||
'level': 'INFO',
|
||||
'project': 'project',
|
||||
'model_name': 'model_name',
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
'type_1',
|
||||
)
|
||||
764
tests/orchestrator/utils/test_orchestrator_functions.py
Normal file
764
tests/orchestrator/utils/test_orchestrator_functions.py
Normal file
@@ -0,0 +1,764 @@
|
||||
from unittest.mock import call, patch
|
||||
|
||||
from orchestrator.utils.orchestrator_functions import (
|
||||
base_scouter,
|
||||
build_tag_config,
|
||||
common_config,
|
||||
drift,
|
||||
gather_read_tags,
|
||||
minimal_retrain,
|
||||
overlap_filter_config,
|
||||
pi_web_api_scouter,
|
||||
predictions_batch,
|
||||
process_path_priority,
|
||||
scouter,
|
||||
simple_metrics,
|
||||
)
|
||||
|
||||
|
||||
def test_common_config():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
}
|
||||
result = common_config(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_common_config_preserves_explicit_runtime():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name'},
|
||||
'runtime': 'tenant-x',
|
||||
}
|
||||
assert common_config(config)['runtime'] == 'tenant-x'
|
||||
|
||||
|
||||
def test_drift():
|
||||
config = {
|
||||
'workflow_type': 'drift',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'interval_minutes': 120,
|
||||
'drift_metrics': ['kolmogorov_smirnov', 'jensen_shannon'],
|
||||
}
|
||||
result = drift(config)
|
||||
expected = {
|
||||
'workflow_type': 'drift',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'schema': 'sientia_data',
|
||||
'source_table_name': 'laborious_data',
|
||||
'target_table_name': 'drift_metrics',
|
||||
'interval': 120,
|
||||
'drift_metrics': ['kolmogorov_smirnov', 'jensen_shannon'],
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_simple_metrics():
|
||||
config = {
|
||||
'workflow_type': 'simple_metrics',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'interval_minutes': 120,
|
||||
'metrics': ['rmse', 'mse'],
|
||||
}
|
||||
result = simple_metrics(config)
|
||||
expected = {
|
||||
'workflow_type': 'simple_metrics',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'schema': 'sientia_data',
|
||||
'predictions_table_name': 'predictions',
|
||||
'data_table_name': 'laborious_data',
|
||||
'target_table_name': 'simple_metrics',
|
||||
'interval_minutes': 120,
|
||||
'metrics': ['rmse', 'mse'],
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_minimal_retrain():
|
||||
config = {
|
||||
'workflow_type': 'minimal_retrain',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;',
|
||||
'datetime_columns': ['timestamp'],
|
||||
}
|
||||
result = minimal_retrain(config)
|
||||
expected = {
|
||||
'workflow_type': 'minimal_retrain',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'query': 'select * from sientia_data.laborious_data order by "timestamp" desc limit 30;',
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_retrain',
|
||||
'datetime_columns': ['timestamp'],
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_scouter():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
|
||||
'read_tags': [
|
||||
{'tag_name': 'test_tag_name', 'aggr_func': 'test_aggr_func', 'data_range': [1, 2]}
|
||||
],
|
||||
'tag_retention_minutes': 10,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
result = scouter(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'topic': 'raw_test_schedule',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter_name': {'policy': 'test_policy'}},
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': 10 * 60,
|
||||
'model_tags': {'test_tag_name': {'aggr_func': 'test_aggr_func', 'data_range': [1, 2]}},
|
||||
'debug_data_package': False,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'fill_missing_tags': False,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_overlap_filter_config():
|
||||
config = [
|
||||
{'filter_name': 'test_filter_name', 'policy': 'test_policy'},
|
||||
{'filter_name': 'test_filter_name2', 'policy': 'test_policy2'},
|
||||
]
|
||||
result = overlap_filter_config({'test_filter_name': {'policy': 'test_policy'}}, config)
|
||||
expected = {
|
||||
'test_filter_name': {'policy': 'test_policy', 'config': {}},
|
||||
'test_filter_name2': {'policy': 'test_policy2', 'config': {}},
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_process_path_priority():
|
||||
config = ['OTHER', 'STOP', 'CONTINUE']
|
||||
result = process_path_priority(config)
|
||||
expected = ['STOP', 'CONTINUE', 'REPEAT']
|
||||
assert result == expected
|
||||
|
||||
|
||||
@patch(
|
||||
'orchestrator.utils.orchestrator_functions.overlap_filter_config',
|
||||
return_value={'test_filter_name': {'policy': 'test_policy', 'config': {}}},
|
||||
)
|
||||
@patch(
|
||||
'orchestrator.utils.orchestrator_functions.process_path_priority',
|
||||
return_value=['STOP', 'CONTINUE', 'REPEAT'],
|
||||
)
|
||||
def test_predictions_batch(mock_process_path_priority, mock_overlap_filter_config):
|
||||
config = {
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_type': 'predictions_batch',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'query': 'test_query',
|
||||
'write_tags': [
|
||||
{'server_id': 'test_server_id', 'type': 'prediction', 'addr': 'test_addr'},
|
||||
{'server_id': 'test_server_id', 'type': 'confidence', 'addr': 'test_addr'},
|
||||
],
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'test_endpoint',
|
||||
'prediction_tags': {
|
||||
'tag_1': 'webid_1',
|
||||
},
|
||||
'confidence_tags': {
|
||||
'tag_2': 'webid_2',
|
||||
},
|
||||
},
|
||||
'input_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
|
||||
'mlflow_transform_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
|
||||
'mlflow_predict_filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'datetime_columns': ['timestamp'],
|
||||
'predictions_storage_policy': 'erl:1',
|
||||
}
|
||||
|
||||
result = predictions_batch(config)
|
||||
|
||||
mock_overlap_filter_config.assert_has_calls(
|
||||
[call({'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, config['input_filters'])]
|
||||
)
|
||||
mock_overlap_filter_config.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
{
|
||||
'EMPTY_DATA': {'policy': 'STOP', 'config': {}},
|
||||
'API_ERROR': {'policy': 'STOP', 'config': {}},
|
||||
},
|
||||
config['mlflow_transform_filters'],
|
||||
)
|
||||
]
|
||||
)
|
||||
mock_overlap_filter_config.assert_has_calls(
|
||||
[call({'API_ERROR': {'policy': 'STOP', 'config': {}}}, config['mlflow_predict_filters'])]
|
||||
)
|
||||
mock_process_path_priority.assert_called_once_with(config['path_priority'])
|
||||
|
||||
expected = {
|
||||
'workflow_type': 'predictions_batch',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'query': 'test_query',
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'predictions',
|
||||
'save_transform': True,
|
||||
'transform_table_name': 'transformed_data',
|
||||
'retention_time': 60 * 60,
|
||||
'opc_output_config': {
|
||||
'test_server_id': {
|
||||
'prediction_tags': {'test_addr': {'data_type': 'float'}},
|
||||
'confidence_tags': {'test_addr': {'data_type': 'float'}},
|
||||
}
|
||||
},
|
||||
'pi_web_api_output_config': {
|
||||
'endpoint': 'test_endpoint',
|
||||
'prediction_tags': {
|
||||
'tag_1': 'webid_1',
|
||||
},
|
||||
'confidence_tags': {
|
||||
'tag_2': 'webid_2',
|
||||
},
|
||||
},
|
||||
'input_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
|
||||
'mlflow_transform_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
|
||||
'mlflow_predict_filters': {'test_filter_name': {'policy': 'test_policy', 'config': {}}},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'datetime_columns': ['timestamp'],
|
||||
'predictions_storage_policy': 'erl:1',
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_gather_read_tags():
|
||||
pipelines = [
|
||||
{
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule2',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address2',
|
||||
},
|
||||
{
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address3',
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
result = gather_read_tags(pipelines)
|
||||
|
||||
expected = {
|
||||
'1:test_tag_address': {
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
'topics': ['raw_test_schedule'],
|
||||
},
|
||||
'2:test_tag_address2': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address2',
|
||||
'topics': ['raw_test_schedule2'],
|
||||
},
|
||||
'2:test_tag_address3': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address3',
|
||||
'topics': ['raw_test_schedule2'],
|
||||
},
|
||||
}
|
||||
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_gather_read_tags_filters_non_scouter_workflows():
|
||||
"""Test that gather_read_tags only processes scouter workflow types and ignores others"""
|
||||
pipelines = [
|
||||
{
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_scouter_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'predictions_batch',
|
||||
'schedule_name': 'test_predictions_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address_predictions',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'minimal_retrain',
|
||||
'schedule_name': 'test_retrain_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '3',
|
||||
'server_name': 'test_server_name3',
|
||||
'tag_address': 'test_tag_address_retrain',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'pi_web_api_scouter',
|
||||
'schedule_name': 'test_pi_web_api_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '4',
|
||||
'server_name': 'test_server_name4',
|
||||
'tag_address': 'test_tag_address_pi_web_api',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'drift',
|
||||
'schedule_name': 'test_drift_schedule',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '5',
|
||||
'server_name': 'test_server_name5',
|
||||
'tag_address': 'test_tag_address_drift',
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_scouter_schedule2',
|
||||
'read_tags': [
|
||||
{
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address2',
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
result = gather_read_tags(pipelines)
|
||||
|
||||
# Only scouter workflow types should be included
|
||||
expected = {
|
||||
'1:test_tag_address': {
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
'topics': ['raw_test_scouter_schedule'],
|
||||
},
|
||||
'1:test_tag_address2': {
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address2',
|
||||
'topics': ['raw_test_scouter_schedule2'],
|
||||
},
|
||||
}
|
||||
|
||||
assert result == expected
|
||||
# Ensure non-scouter pipelines are not included
|
||||
assert '2:test_tag_address_predictions' not in result
|
||||
assert '3:test_tag_address_retrain' not in result
|
||||
assert '4:test_tag_address_pi_web_api' not in result
|
||||
assert '5:test_tag_address_drift' not in result
|
||||
|
||||
|
||||
def test_build_tag_config():
|
||||
tags = [
|
||||
{
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
'frequency': 1000,
|
||||
},
|
||||
{
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address2',
|
||||
'frequency': 1000,
|
||||
},
|
||||
{
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address3',
|
||||
'frequency': 300,
|
||||
},
|
||||
{
|
||||
'server_id': '3',
|
||||
'server_name': 'test_server_name3',
|
||||
'tag_address': 'test_tag_address4',
|
||||
'frequency': 200,
|
||||
},
|
||||
]
|
||||
|
||||
opc_servers = {
|
||||
'1': {
|
||||
'server_name': 'test_server_name',
|
||||
'url': 'test_url',
|
||||
'uri': 'test_uri',
|
||||
},
|
||||
'2': {
|
||||
'server_name': 'test_server_name2',
|
||||
'url': 'test_url2',
|
||||
'uri': 'test_uri2',
|
||||
},
|
||||
}
|
||||
|
||||
result = build_tag_config(tags, opc_servers)
|
||||
|
||||
expected = {
|
||||
'test_server_name': {
|
||||
'server_id': '1',
|
||||
'name': 'test_server_name',
|
||||
'url': 'test_url',
|
||||
'server_uri': 'test_uri',
|
||||
'cert_path': None,
|
||||
'private_key_path': None,
|
||||
'server_cert_path': None,
|
||||
'subscription_period_ms': 500,
|
||||
'tags': {
|
||||
'test_tag_address': {
|
||||
'server_id': '1',
|
||||
'server_name': 'test_server_name',
|
||||
'tag_address': 'test_tag_address',
|
||||
'frequency': 1000,
|
||||
},
|
||||
},
|
||||
},
|
||||
'test_server_name2': {
|
||||
'server_id': '2',
|
||||
'name': 'test_server_name2',
|
||||
'url': 'test_url2',
|
||||
'server_uri': 'test_uri2',
|
||||
'cert_path': None,
|
||||
'private_key_path': None,
|
||||
'server_cert_path': None,
|
||||
'subscription_period_ms': 150,
|
||||
'tags': {
|
||||
'test_tag_address2': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address2',
|
||||
'frequency': 1000,
|
||||
},
|
||||
'test_tag_address3': {
|
||||
'server_id': '2',
|
||||
'server_name': 'test_server_name2',
|
||||
'tag_address': 'test_tag_address3',
|
||||
'frequency': 300,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
assert result == (expected, ['3'])
|
||||
|
||||
|
||||
def test_base_scouter():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'filters': [
|
||||
{'filter_name': 'test_filter_name', 'policy': 'test_policy'},
|
||||
{'filter_name': 'test_filter_name2', 'policy': 'test_policy2'},
|
||||
],
|
||||
'tag_retention_minutes': 30,
|
||||
'debug_data_package': True,
|
||||
'fill_missing_tags': True,
|
||||
}
|
||||
result = base_scouter(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'trigger_laborious': False,
|
||||
'filters': {
|
||||
'test_filter_name': {'policy': 'test_policy'},
|
||||
'test_filter_name2': {'policy': 'test_policy2'},
|
||||
},
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': 30 * 60,
|
||||
'debug_data_package': True,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'fill_missing_tags': True,
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_pi_web_api_scouter():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'filters': [{'filter_name': 'test_filter_name', 'policy': 'test_policy'}],
|
||||
'read_tags': {
|
||||
'test_tag_name': {
|
||||
'webid': 'test_webid',
|
||||
'aggr_func': 'test_aggr_func',
|
||||
'data_range': [1, 2],
|
||||
}
|
||||
},
|
||||
'tag_retention_minutes': 10,
|
||||
'pi_web_api_config': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
'period': '*-2d',
|
||||
'max_count': 5,
|
||||
'api_timeout': 30,
|
||||
},
|
||||
'frequency': '1m',
|
||||
}
|
||||
result = pi_web_api_scouter(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter_name': {'policy': 'test_policy'}},
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': 10 * 60,
|
||||
'model_tags': {
|
||||
'test_tag_name': {
|
||||
'webid': 'test_webid',
|
||||
'aggr_func': 'test_aggr_func',
|
||||
'data_range': [1, 2],
|
||||
}
|
||||
},
|
||||
'debug_data_package': False,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'fill_missing_tags': False,
|
||||
'pi_web_api_query': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
'period': '*-2d',
|
||||
'max_count': 5,
|
||||
'api_timeout': 30,
|
||||
},
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_pi_web_api_scouter_with_timeout_greater_than_frequency():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'filters': [],
|
||||
'read_tags': {
|
||||
'test_tag_name': {
|
||||
'webid': 'test_webid',
|
||||
}
|
||||
},
|
||||
'tag_retention_minutes': 10,
|
||||
'pi_web_api_config': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
'api_timeout': 120,
|
||||
},
|
||||
'frequency': '1m',
|
||||
}
|
||||
result = pi_web_api_scouter(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '1m',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'trigger_laborious': False,
|
||||
'filters': {},
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': 10 * 60,
|
||||
'model_tags': {
|
||||
'test_tag_name': {'webid': 'test_webid', 'aggr_func': 'lts', 'data_range': [-100, 100]}
|
||||
},
|
||||
'debug_data_package': False,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'fill_missing_tags': False,
|
||||
'pi_web_api_query': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
'period': '*-1d',
|
||||
'max_count': 1,
|
||||
'api_timeout': 60,
|
||||
},
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
|
||||
|
||||
def test_pi_web_api_scouter_with_no_timeout():
|
||||
config = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_id': 'test_model_id',
|
||||
'model': {'name': 'test_model_name', 'model_config': {'test_config': 'test_config'}},
|
||||
'filters': [],
|
||||
'read_tags': {
|
||||
'test_tag_name': {
|
||||
'webid': 'test_webid',
|
||||
}
|
||||
},
|
||||
'tag_retention_minutes': 10,
|
||||
'pi_web_api_config': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
},
|
||||
'frequency': '30s',
|
||||
}
|
||||
result = pi_web_api_scouter(config)
|
||||
expected = {
|
||||
'workflow_type': 'scouter',
|
||||
'schedule_name': 'test_schedule',
|
||||
'frequency': '30s',
|
||||
'offset': '0m',
|
||||
'max_retry_policy': 1,
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model_name',
|
||||
'model_config': {'test_config': 'test_config'},
|
||||
'on_conflict': 'error',
|
||||
'trigger_laborious': False,
|
||||
'filters': {},
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'laborious_data',
|
||||
'retention_time': 10 * 60,
|
||||
'model_tags': {
|
||||
'test_tag_name': {'webid': 'test_webid', 'aggr_func': 'lts', 'data_range': [-100, 100]}
|
||||
},
|
||||
'debug_data_package': False,
|
||||
'execution_timeout_seconds': 300,
|
||||
'task_timeout_seconds': 300,
|
||||
'fill_missing_tags': False,
|
||||
'pi_web_api_query': {
|
||||
'endpoint': 'https://test-endpoint.com',
|
||||
'period': '*-1d',
|
||||
'max_count': 1,
|
||||
'api_timeout': 30,
|
||||
},
|
||||
'runtime': 'legacy',
|
||||
}
|
||||
assert result == expected
|
||||
49
tests/orchestrator/worker/test_worker.py
Normal file
49
tests/orchestrator/worker/test_worker.py
Normal file
@@ -0,0 +1,49 @@
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from orchestrator.worker.worker import main
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
from orchestrator.workflows.reports import Reports
|
||||
|
||||
|
||||
def test_main_is_coroutine():
|
||||
assert asyncio.iscoroutinefunction(main)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.worker.worker.sys')
|
||||
@patch('orchestrator.worker.worker.start_http_server')
|
||||
@patch('orchestrator.worker.worker.NotificationHandler')
|
||||
@patch('orchestrator.worker.worker.Activities')
|
||||
@patch('orchestrator.worker.worker.prepare_worker')
|
||||
@patch('orchestrator.worker.worker.client.Client.connect', new_callable=AsyncMock)
|
||||
async def test_main_starts_three_workers_for_expected_workflows(
|
||||
_connect_mock,
|
||||
prepare_worker_mock,
|
||||
activities_mock,
|
||||
_notification_handler_mock,
|
||||
_start_http_server,
|
||||
_sys_mock,
|
||||
):
|
||||
"""
|
||||
main() must spin up exactly three Temporal workers, one per main workflow
|
||||
(Orchestrator, Alerts, Reports), and call run() on each.
|
||||
"""
|
||||
activities_instance = activities_mock.return_value
|
||||
activities_instance.connect_to_temporal = AsyncMock()
|
||||
|
||||
worker_mock = MagicMock()
|
||||
worker_mock.run = AsyncMock()
|
||||
prepare_worker_mock.return_value = worker_mock
|
||||
|
||||
await main()
|
||||
|
||||
assert prepare_worker_mock.call_count == 3
|
||||
|
||||
main_workflows = [call.kwargs['main_workflow'] for call in prepare_worker_mock.call_args_list]
|
||||
assert main_workflows == [Orchestrator, Alerts, Reports]
|
||||
|
||||
assert worker_mock.run.call_count == 3
|
||||
0
tests/orchestrator/workflows/__init__.py
Normal file
0
tests/orchestrator/workflows/__init__.py
Normal file
@@ -0,0 +1,171 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.workflows.subworkflows.load_notification_package import LoadNotificationPackage
|
||||
|
||||
|
||||
@fixture
|
||||
def load_notification_package():
|
||||
return LoadNotificationPackage()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'workflow_name': 'test-workflow',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
|
||||
)
|
||||
async def test_run(workflow_mock, load_notification_package):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'mail_type': 'test_mail_type',
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
}
|
||||
|
||||
workflow_mock.start_local_activity_method.side_effect = [
|
||||
'2023-01-01 12:00:00',
|
||||
[
|
||||
{
|
||||
'id': '1',
|
||||
}
|
||||
],
|
||||
[
|
||||
{
|
||||
'id_r': '1',
|
||||
}
|
||||
],
|
||||
]
|
||||
|
||||
output = await load_notification_package.run(input_data)
|
||||
|
||||
assert output == {
|
||||
'last_timestamp': '2023-01-01 12:00:00',
|
||||
'notification_package': [
|
||||
{
|
||||
'id': '1',
|
||||
}
|
||||
],
|
||||
'sending_configs': [
|
||||
{
|
||||
'id_r': '1',
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.get_last_data_timestamp,
|
||||
{**input_data['metadata'], 'mail_type': 'test_mail_type'},
|
||||
start_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{
|
||||
**input_data['metadata'],
|
||||
'query': {'collection': 'receiver_groups', 'filters': {'active': True}},
|
||||
},
|
||||
start_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_latest_data,
|
||||
{
|
||||
**input_data['metadata'],
|
||||
'collection_name': 'notification_queue',
|
||||
'last_data_timestamp': '2023-01-01 12:00:00',
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
},
|
||||
start_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.put_last_data_timestamp,
|
||||
{
|
||||
**input_data['metadata'],
|
||||
'data': [
|
||||
{
|
||||
'id': '1',
|
||||
}
|
||||
],
|
||||
'mail_type': 'test_mail_type',
|
||||
},
|
||||
start_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
|
||||
)
|
||||
async def test_run_no_data(workflow_mock, load_notification_package):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'mail_type': 'test_mail_type',
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
}
|
||||
|
||||
workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', [], []]
|
||||
|
||||
output = await load_notification_package.run(input_data)
|
||||
|
||||
assert output == {
|
||||
'last_timestamp': '2023-01-01 12:00:00',
|
||||
'notification_package': [],
|
||||
'sending_configs': [],
|
||||
}
|
||||
|
||||
workflow_mock.start_activity_method.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch(
|
||||
'orchestrator.workflows.subworkflows.load_notification_package.workflow', new_callable=AsyncMock
|
||||
)
|
||||
async def test_run_has_data(workflow_mock, load_notification_package):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'base_data_filter': {'level': 'ERROR'},
|
||||
'mail_type': 'test_mail_type',
|
||||
}
|
||||
|
||||
workflow_mock.start_local_activity_method.side_effect = ['2023-01-01 12:00:00', ['data'], []]
|
||||
|
||||
output = await load_notification_package.run(input_data)
|
||||
|
||||
assert output == {
|
||||
'last_timestamp': '2023-01-01 12:00:00',
|
||||
'notification_package': ['data'],
|
||||
'sending_configs': [],
|
||||
}
|
||||
|
||||
workflow_mock.start_activity_method.assert_not_called()
|
||||
@@ -0,0 +1,151 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.workflows.subworkflows.process_notifications import ProcessNotifications
|
||||
|
||||
|
||||
@fixture
|
||||
def process_notifications():
|
||||
return ProcessNotifications()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'workflow_name': 'test-workflow',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.subworkflows.process_notifications.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock, process_notifications):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'notification_package': ['content'],
|
||||
'mail_type': 'test_mail_type',
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table_name',
|
||||
}
|
||||
|
||||
response = await process_notifications.run(input_data)
|
||||
|
||||
assert response == workflow_mock.execute_local_activity_method.return_value
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.build_email_html,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': input_data['notification_package'],
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.format_log_report,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': workflow_mock.execute_activity_method.return_value,
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.send_email,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.subworkflows.process_notifications.workflow', new_callable=AsyncMock)
|
||||
async def test_run_send_email_return_empty(workflow_mock, process_notifications):
|
||||
input_data = {
|
||||
'metadata': metadata,
|
||||
'notification_package': ['content'],
|
||||
'mail_type': 'test_mail_type',
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table_name',
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method.return_value = []
|
||||
|
||||
await process_notifications.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.build_email_html,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': input_data['notification_package'],
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.send_email,
|
||||
{
|
||||
**metadata,
|
||||
'receiver_groups': workflow_mock.execute_local_activity_method.return_value,
|
||||
'mail_type': input_data['mail_type'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert workflow_mock.execute_activity_method.call_count == 1
|
||||
assert workflow_mock.execute_local_activity_method.call_count == 1
|
||||
137
tests/orchestrator/workflows/test_alerts.py
Normal file
137
tests/orchestrator/workflows/test_alerts.py
Normal file
@@ -0,0 +1,137 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.workflows.alerts import Alerts
|
||||
|
||||
|
||||
@fixture
|
||||
def alerts():
|
||||
return Alerts()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'workflow_name': 'alerts',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
|
||||
async def test_run_full_flow(workflow_mock, alerts):
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await alerts.run(input_data)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.load_notification_package',
|
||||
{**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.process_notifications',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'mail_type': 'Alerts',
|
||||
'notification_package': workflow_mock.execute_local_activity_method.return_value,
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_report',
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.filter_notification_alerts,
|
||||
{
|
||||
**metadata,
|
||||
'notification_package': workflow_mock.execute_child_workflow.return_value[
|
||||
'notification_package'
|
||||
],
|
||||
'sending_configs': workflow_mock.execute_child_workflow.return_value[
|
||||
'sending_configs'
|
||||
],
|
||||
'notification_ttl': input_data['notification_ttl'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.store_notification_cache,
|
||||
{
|
||||
**metadata,
|
||||
'log_report': workflow_mock.execute_child_workflow.return_value,
|
||||
'sent_ttl': input_data['sent_ttl'],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
|
||||
async def test_run_no_data(workflow_mock, alerts):
|
||||
workflow_mock.execute_child_workflow.return_value = {
|
||||
'last_timestamp': '2023-01-01 12:00:00.000000',
|
||||
'notification_package': [],
|
||||
'sending_configs': [],
|
||||
}
|
||||
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await alerts.run(input_data)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.load_notification_package',
|
||||
{**input_data, 'metadata': metadata, 'base_data_filter': {'level': 'ERROR'}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
|
||||
async def test_run_no_receiver_groups(workflow_mock, alerts):
|
||||
workflow_mock.execute_local_activity_method.return_value = {}
|
||||
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await alerts.run(input_data)
|
||||
|
||||
assert workflow_mock.execute_child_workflow.call_count == 1
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.alerts.workflow', new_callable=AsyncMock)
|
||||
async def test_run_no_log_report(workflow_mock, alerts):
|
||||
workflow_mock.execute_child_workflow.side_effect = [MagicMock(), []]
|
||||
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await alerts.run(input_data)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_not_called()
|
||||
348
tests/orchestrator/workflows/test_orchestrator.py
Normal file
348
tests/orchestrator/workflows/test_orchestrator.py
Normal file
@@ -0,0 +1,348 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.workflows.orchestrator import Orchestrator
|
||||
|
||||
|
||||
@fixture
|
||||
def orchestrator():
|
||||
return Orchestrator()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'workflow_name': 'orchestrator',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.orchestrator.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock, orchestrator):
|
||||
input_data = {
|
||||
'pipelines_query': 'SELECT * FROM bucket',
|
||||
'opc_servers_query': 'SELECT * FROM servers',
|
||||
'schedule_name': 'test-schedule-name',
|
||||
}
|
||||
|
||||
await orchestrator.run(input_data)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.aggregate_documents_in_mongodb,
|
||||
{
|
||||
**metadata,
|
||||
'query': input_data['pipelines_query'],
|
||||
'timestamp_fields': ['updated_at'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{**metadata, 'query': input_data['opc_servers_query']},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.find_documents_in_mongodb,
|
||||
{
|
||||
**metadata,
|
||||
'query': {'collection': 'orchestrated_schedules'},
|
||||
'timestamp_fields': ['updated_at'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_opc_slots,
|
||||
{**metadata},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_active_ingestors,
|
||||
{**metadata},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.format_schedule_config,
|
||||
{
|
||||
**metadata,
|
||||
'schedule_config': workflow_mock.start_local_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.process_schedules,
|
||||
{**metadata, 'pipelines': workflow_mock.start_local_activity_method.return_value},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.process_slots,
|
||||
{
|
||||
**metadata,
|
||||
'opc_servers': workflow_mock.start_local_activity_method.return_value,
|
||||
'active_ingestors': workflow_mock.start_local_activity_method.return_value,
|
||||
'pipelines': workflow_mock.start_local_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.create_schedule_config,
|
||||
{
|
||||
**metadata,
|
||||
'current_schedule_config': workflow_mock.start_local_activity_method.return_value,
|
||||
'schedule_config': workflow_mock.start_local_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.create_slot_config,
|
||||
{
|
||||
**metadata,
|
||||
'current_slot_config': workflow_mock.start_local_activity_method.return_value,
|
||||
'slot_config': workflow_mock.start_local_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.normalize_schedules,
|
||||
{
|
||||
**metadata,
|
||||
'orchestrated_schedules': workflow_mock.start_local_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.create_collection_with_ttl_index,
|
||||
{
|
||||
**metadata,
|
||||
'pipelines': workflow_mock.start_local_activity_method.return_value['scouter'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.delete_slots,
|
||||
{
|
||||
**metadata,
|
||||
'to_delete': workflow_mock.start_local_activity_method.return_value[
|
||||
'to_delete'
|
||||
],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.update_slots,
|
||||
{
|
||||
**metadata,
|
||||
'to_insert': workflow_mock.start_local_activity_method.return_value[
|
||||
'to_insert'
|
||||
],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.delete_schedules,
|
||||
{
|
||||
**metadata,
|
||||
'schedules': workflow_mock.start_local_activity_method.return_value[
|
||||
'to_delete'
|
||||
],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.create_schedules,
|
||||
{
|
||||
**metadata,
|
||||
'schedules': workflow_mock.start_local_activity_method.return_value[
|
||||
'to_create'
|
||||
],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.update_schedules,
|
||||
{
|
||||
**metadata,
|
||||
'schedules': workflow_mock.start_local_activity_method.return_value[
|
||||
'to_update'
|
||||
],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.report_schedule_orchestration,
|
||||
{
|
||||
**metadata,
|
||||
'created_schedules': workflow_mock.start_activity_method.return_value,
|
||||
'updated_schedules': workflow_mock.start_activity_method.return_value,
|
||||
'deleted_schedules': workflow_mock.start_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.report_slot_orchestration,
|
||||
{
|
||||
**metadata,
|
||||
'inserted_slots': workflow_mock.start_activity_method.return_value,
|
||||
'deleted_slots': workflow_mock.start_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.update_pipelines_timestamps,
|
||||
{
|
||||
**metadata,
|
||||
'updated_pipelines': workflow_mock.start_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.delete_pipelines_timestamps,
|
||||
{
|
||||
**metadata,
|
||||
'deleted_pipelines': workflow_mock.start_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.start_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.create_pipelines_timestamps,
|
||||
{
|
||||
**metadata,
|
||||
'created_pipelines': workflow_mock.start_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
109
tests/orchestrator/workflows/test_reports.py
Normal file
109
tests/orchestrator/workflows/test_reports.py
Normal file
@@ -0,0 +1,109 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
from pytest import fixture, mark
|
||||
|
||||
from orchestrator.activities.activities import Activities
|
||||
from orchestrator.workflows.reports import Reports
|
||||
|
||||
|
||||
@fixture
|
||||
def reports():
|
||||
return Reports()
|
||||
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule-name',
|
||||
'workflow_name': 'reports',
|
||||
'model_name': '-',
|
||||
'model_id': '-',
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
|
||||
async def test_run_full_flow(workflow_mock, reports):
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await reports.run(input_data)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.load_notification_package',
|
||||
{**input_data, 'metadata': metadata, 'base_data_filter': {}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.process_notifications',
|
||||
{
|
||||
'metadata': metadata,
|
||||
'mail_type': 'Reports',
|
||||
'notification_package': workflow_mock.execute_local_activity_method.return_value,
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'log_report',
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.filter_notification_reports,
|
||||
{
|
||||
**metadata,
|
||||
'notification_package': workflow_mock.execute_child_workflow.return_value[
|
||||
'notification_package'
|
||||
],
|
||||
'sending_configs': workflow_mock.execute_child_workflow.return_value[
|
||||
'sending_configs'
|
||||
],
|
||||
},
|
||||
schedule_to_close_timeout=ANY,
|
||||
retry_policy=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
|
||||
async def test_run_no_data(workflow_mock, reports):
|
||||
workflow_mock.execute_child_workflow.return_value = {
|
||||
'last_timestamp': '2023-01-01 12:00:00.000000',
|
||||
'notification_package': [],
|
||||
'sending_configs': [],
|
||||
}
|
||||
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await reports.run(input_data)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
'subworkflow.load_notification_package',
|
||||
{**input_data, 'metadata': metadata, 'base_data_filter': {}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('orchestrator.workflows.reports.workflow', new_callable=AsyncMock)
|
||||
async def test_run_no_groups(workflow_mock, reports):
|
||||
workflow_mock.execute_local_activity_method.return_value = []
|
||||
|
||||
input_data = {'schedule_name': 'test-schedule-name', 'notification_ttl': 300, 'sent_ttl': 600}
|
||||
|
||||
await reports.run(input_data)
|
||||
|
||||
workflow_mock.execute_child_workflow.assert_called_once()
|
||||
Reference in New Issue
Block a user