SIENTIAPDE-1325

Remove unused Couchbase configurations and related code. Refactor Activities class to eliminate Couchbase dependency, updating initialization and shutdown methods. Update README and tests to reflect these changes. Upgrade sientia-dataops-library dependency version in requirements.txt.
This commit is contained in:
vitor-aignosi
2025-11-06 10:45:23 -03:00
parent f099366bed
commit acde9475ec
20 changed files with 458 additions and 614 deletions

View File

@@ -3,10 +3,6 @@ REDIS_PORT="6379"
REDIS_USERNAME="redis_username" REDIS_USERNAME="redis_username"
REDIS_PASSWORD="redis_password" REDIS_PASSWORD="redis_password"
COUCHBASE_CONNECTION_STRING="couchbase://sientia.couchbase.svc.cluster.local"
COUCHBASE_USERNAME="sientia"
COUCHBASE_PASSWORD="sientia"
MONGODB_USERNAME="mongo_username" MONGODB_USERNAME="mongo_username"
MONGODB_PASSWORD="mongo_password" MONGODB_PASSWORD="mongo_password"
MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017" MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"

1
.gitignore vendored
View File

@@ -13,7 +13,6 @@ docker-compose.override.yml
scouter/.file_versions/ scouter/.file_versions/
scouter/pipelines/**/triggers.yaml scouter/pipelines/**/triggers.yaml
**/postgres_data/** **/postgres_data/**
**/couchbase_data/**
**/redis_data/** **/redis_data/**
# Ignorar arquivos e diretórios de cache do Python # Ignorar arquivos e diretórios de cache do Python
__pycache__/ __pycache__/

View File

@@ -408,7 +408,6 @@ The orchestrator includes an advanced notification filtering system that prevent
- **Email**: SMTP operations with HTML generation and attachment support - **Email**: SMTP operations with HTML generation and attachment support
- **Formatters**: Configuration processing, slot distribution algorithms, and notification filtering for reports - **Formatters**: Configuration processing, slot distribution algorithms, and notification filtering for reports
- **Postgres**: PostgreSQL operations for audit logging and data export (via sientia-dataops-library) - **Postgres**: PostgreSQL operations for audit logging and data export (via sientia-dataops-library)
- **Couchbase**: Database operations (currently unused but maintained for future use)
#### **Utilities (`orchestrator/utils/`)** #### **Utilities (`orchestrator/utils/`)**
- **Connectors Configuration**: Database and service configuration management - **Connectors Configuration**: Database and service configuration management
@@ -642,7 +641,6 @@ The Orchestrator system exposes comprehensive Prometheus metrics:
tests/ tests/
├── orchestrator/ # Orchestrator workflow tests ├── orchestrator/ # Orchestrator workflow tests
│ ├── test_activities.py │ ├── test_activities.py
│ ├── test_couchbase.py
│ ├── test_email.py │ ├── test_email.py
│ ├── test_formatters.py │ ├── test_formatters.py
│ ├── test_mongo_db.py │ ├── test_mongo_db.py
@@ -755,7 +753,6 @@ orchestrator/
│ ├── mongo_db.py # MongoDB operations │ ├── mongo_db.py # MongoDB operations
│ ├── email.py # Email service operations │ ├── email.py # Email service operations
│ ├── formatters.py # Configuration formatting and report filtering │ ├── formatters.py # Configuration formatting and report filtering
│ └── couchbase.py # Couchbase operations (currently unused)
├── workflows/ # Temporal workflow definitions ├── workflows/ # Temporal workflow definitions
│ ├── orchestrator.py # Main orchestration workflow │ ├── orchestrator.py # Main orchestration workflow
│ ├── alerts.py # Error alert workflow │ ├── alerts.py # Error alert workflow

View File

@@ -5,6 +5,7 @@ with workflow.unsafe.imports_passed_through():
from typing import Any from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.temporal.activities.postgres import Postgres from sientia_do.temporal.activities.postgres import Postgres
from orchestrator.activities.email import Email from orchestrator.activities.email import Email
@@ -14,9 +15,7 @@ with workflow.unsafe.imports_passed_through():
from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.activities.temporal_manager import TemporalManager
class Activities( # Couchbase, class Activities(TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres):
TemporalManager, SlotManager, Formatters, MongoDB, Email, Postgres
):
""" """
Central activities orchestrator for Temporal workflow operations. Central activities orchestrator for Temporal workflow operations.
@@ -47,11 +46,8 @@ class Activities( # Couchbase,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
): ):
# Initialize parent classes # Initialize parent classes
# Couchbase.__init__(self, connection_string=couchbase_config['connection_string'],
# username=couchbase_config['username'], metrics_controller = MetricsController(logger=logger)
# password=couchbase_config['password'],
# logger=logger,
# notification_handler=notification_handler)
TemporalManager.__init__( TemporalManager.__init__(
self, self,
@@ -60,6 +56,7 @@ class Activities( # Couchbase,
laborious_namespace=temporal_config['temporal_laborious_namespace'], laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
SlotManager.__init__( SlotManager.__init__(
@@ -70,6 +67,7 @@ class Activities( # Couchbase,
password=redis_config['password'], password=redis_config['password'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
Formatters.__init__( Formatters.__init__(
@@ -78,6 +76,7 @@ class Activities( # Couchbase,
laborious_namespace=temporal_config['temporal_laborious_namespace'], laborious_namespace=temporal_config['temporal_laborious_namespace'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
MongoDB.__init__( MongoDB.__init__(
@@ -87,6 +86,7 @@ class Activities( # Couchbase,
ttl_index_seconds=mongodb_config['ttl_index_seconds'], ttl_index_seconds=mongodb_config['ttl_index_seconds'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
Email.__init__( Email.__init__(
@@ -97,6 +97,7 @@ class Activities( # Couchbase,
smtp_port=email_config['smtp_port'], smtp_port=email_config['smtp_port'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
Postgres.__init__( Postgres.__init__(
@@ -110,6 +111,7 @@ class Activities( # Couchbase,
max_connections=postgres_config['max_connections'], max_connections=postgres_config['max_connections'],
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
def shutdown(self): def shutdown(self):
@@ -120,6 +122,9 @@ class Activities( # Couchbase,
services, and other resources to ensure proper cleanup when the services, and other resources to ensure proper cleanup when the
application terminates. application terminates.
""" """
MongoDB.shutdown(self) MongoDB.close(self)
Postgres.close(self) Postgres.close(self)
Email.shutdown(self) Email.close(self)
Formatters.close(self)
SlotManager.close(self)
TemporalManager.close(self)

View File

@@ -1,112 +0,0 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import json
import traceback
from datetime import timedelta
from logging import Logger
from typing import Any
from couchbase.auth import PasswordAuthenticator
from couchbase.cluster import Cluster
from couchbase.options import ClusterOptions
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity
class Couchbase(BaseActivity):
"""
Couchbase database operations activity (currently unused).
This class provides Couchbase database connectivity and query operations
for Temporal workflows. It handles connection management, query execution,
and error reporting with automatic connection lifecycle management.
Args:
connection_string (str): Couchbase cluster connection string
username (str): Couchbase authentication username
password (str): Couchbase authentication password
logger (Logger): Application logger instance
notification_handler (NotificationHandler): Notification management handler
Note:
This class is currently commented out in the main Activities class
but maintained for potential future use.
"""
def __init__(
self,
connection_string: str,
username: str,
password: str,
logger: Logger,
notification_handler: NotificationHandler,
):
self.connection_string = connection_string
self.username = username
self.password = password
logger.info('Initializing Couchbase connection...')
self.cluster = Cluster(
connection_string,
ClusterOptions(
authenticator=PasswordAuthenticator(username=username, password=password)
),
)
logger.info('Awaiting Couchbase connection...')
self.cluster.wait_until_ready(timeout=timedelta(seconds=10))
logger.info('Couchbase connection ready')
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler)
def shutdown(self):
try:
self.cluster.close()
except Exception as e:
self.logger.error(f'Failed to close Couchbase connection: {e}')
def __del__(self):
self.shutdown()
@activity.defn(name='load_query_from_couchbase')
async def load_query_from_couchbase(self, input_data: dict[str, Any]) -> list[dict[str, Any]]:
"""
Load a query from couchbase
Args:
input_data (dict[str, Any]): The input data containing the query to execute
Returns:
list[dict[str, Any]]: The result of the query
"""
query = input_data['query']
self.logger.info(f'Executing couchbase query: {query}')
try:
result = self.cluster.query(query)
except Exception as e:
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id='COUCHBASE_LOAD_QUERY_ERROR',
message=f'Failed to execute couchbase query: {e}',
block='load_query_from_couchbase',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.logger.error(trace)
raise e
rows = []
for row in result.rows():
rows.append(row)
self.logger.info('Fetched %d rows from couchbase', len(rows))
self.logger.debug('Rows: \n %s', json.dumps(rows, indent=4, sort_keys=True))
return rows

View File

@@ -12,13 +12,13 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.observability.logger import Logger from sientia_do.observability.logger import Logger
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from orchestrator import metrics from orchestrator import metrics
from orchestrator.utils.email_builder import EmailBuilder from orchestrator.utils.email_builder import EmailBuilder
class Email(BaseActivity): class Email(SientiaMonitoring):
""" """
Email service activity for sending workflow notifications. Email service activity for sending workflow notifications.
@@ -43,6 +43,7 @@ class Email(BaseActivity):
smtp_port: int, smtp_port: int,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController,
): ):
self.email_builder = EmailBuilder(logger=logger) self.email_builder = EmailBuilder(logger=logger)
@@ -60,13 +61,25 @@ class Email(BaseActivity):
self.server.starttls() self.server.starttls()
self.server.login(self.sender_email, self.sender_password) self.server.login(self.sender_email, self.sender_password)
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
def shutdown(self): def close(self):
""" """
Shutdown the Email connection and clean up resources. Close the Email connection and clean up resources.
""" """
self.server.quit() 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') @activity.defn(name='build_email_html')
async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]: async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:

View File

@@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through():
from pandas import DataFrame from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
from orchestrator.utils.orchestrator_functions import ( from orchestrator.utils.orchestrator_functions import (
@@ -24,7 +24,7 @@ with workflow.unsafe.imports_passed_through():
topic_separator = '\n ========== \n' topic_separator = '\n ========== \n'
class Formatters(BaseActivity): class Formatters(SientiaMonitoring):
""" """
Schedule and slot configuration formatting and notification filtering activity. Schedule and slot configuration formatting and notification filtering activity.
@@ -53,10 +53,28 @@ class Formatters(BaseActivity):
laborious_namespace: str, laborious_namespace: str,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController,
): ):
self.scouter_namespace = scouter_namespace self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace self.laborious_namespace = laborious_namespace
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
def close(self):
"""
Close the Formatters connection and clean up 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') @activity.defn(name='process_schedules')
async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]: async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -166,7 +184,7 @@ class Formatters(BaseActivity):
slot_config[f'{i}'], notifications = build_tag_config(slot_tags, opc_servers) slot_config[f'{i}'], notifications = build_tag_config(slot_tags, opc_servers)
if notifications: if notifications:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION', notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
message=f'Servers {", ".join(notifications)} not found in opc_servers', message=f'Servers {", ".join(notifications)} not found in opc_servers',
@@ -179,7 +197,7 @@ class Formatters(BaseActivity):
slot_tags = tags[last_index:] slot_tags = tags[last_index:]
slot_config[f'{number_of_slots}'], notifications = build_tag_config(slot_tags, opc_servers) slot_config[f'{number_of_slots}'], notifications = build_tag_config(slot_tags, opc_servers)
if notifications: if notifications:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION', notification_id='ORCHESTRATOR_SERVER_NOT_FOUND_DURING_SLOT_CONFIGURATION',
message=f'Servers {", ".join(notifications)} not found in opc_servers', message=f'Servers {", ".join(notifications)} not found in opc_servers',
@@ -357,7 +375,7 @@ class Formatters(BaseActivity):
return output return output
def send_success_report( async def send_success_report(
self, self,
metadata: dict[str, Any], metadata: dict[str, Any],
message: str, message: str,
@@ -373,7 +391,7 @@ class Formatters(BaseActivity):
notification_id (str): The ID of the notification to send. notification_id (str): The ID of the notification to send.
attachment (Any | None, optional): Optional attachment content to include. attachment (Any | None, optional): Optional attachment content to include.
""" """
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id=notification_id, notification_id=notification_id,
message=message, message=message,
@@ -382,7 +400,7 @@ class Formatters(BaseActivity):
attachment_content=json.dumps(attachment, indent=4, sort_keys=True), attachment_content=json.dumps(attachment, indent=4, sort_keys=True),
) )
def send_error_report( async def send_error_report(
self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str self, metadata: dict[str, Any], message: str, notification_id: str, attachment: str
) -> None: ) -> None:
""" """
@@ -394,7 +412,7 @@ class Formatters(BaseActivity):
notification_id (str): The ID of the notification. notification_id (str): The ID of the notification.
attachment (str): The attachment content for the notification. attachment (str): The attachment content for the notification.
""" """
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id=notification_id, notification_id=notification_id,
message=message, message=message,
@@ -455,7 +473,7 @@ class Formatters(BaseActivity):
return success_keys, error_keys return success_keys, error_keys
def manage_and_send_report( async def manage_and_send_report(
self, self,
metadata: dict[str, Any], metadata: dict[str, Any],
success_keys: list[str], success_keys: list[str],
@@ -474,7 +492,7 @@ class Formatters(BaseActivity):
schedule_data (dict[str, Any]): The schedule data containing items and notification ID. schedule_data (dict[str, Any]): The schedule data containing items and notification ID.
""" """
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( await self.send_success_report(
metadata=metadata, metadata=metadata,
message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}', message=f'Successfully {schedule_type}: \n {", ".join(success_keys)}',
notification_id=schedule_data['id'], notification_id=schedule_data['id'],
@@ -489,7 +507,7 @@ class Formatters(BaseActivity):
else: else:
attachment.append(f'{key}:\n{value["message"]}') attachment.append(f'{key}:\n{value["message"]}')
self.send_error_report( await self.send_error_report(
metadata=metadata, metadata=metadata,
message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}', message=f'Fails on {schedule_type}: \n {", ".join(error_keys)}',
notification_id=f'{schedule_data["id"]}_ERROR', notification_id=f'{schedule_data["id"]}_ERROR',
@@ -535,7 +553,7 @@ class Formatters(BaseActivity):
if len(schedule_data['items']) > 0: if len(schedule_data['items']) > 0:
success_keys, error_keys = self.parse_report_schedule(schedule_data['items']) success_keys, error_keys = self.parse_report_schedule(schedule_data['items'])
self.manage_and_send_report( await self.manage_and_send_report(
metadata=metadata, metadata=metadata,
success_keys=success_keys, success_keys=success_keys,
error_keys=error_keys, error_keys=error_keys,
@@ -566,14 +584,14 @@ class Formatters(BaseActivity):
success_keys, error_keys = self.parse_report(inserted_slots) success_keys, error_keys = self.parse_report(inserted_slots)
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( await self.send_success_report(
metadata=metadata, metadata=metadata,
message=f'Inserted slots: \n {", ".join(success_keys)}', message=f'Inserted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
) )
if len(error_keys) > 0: if len(error_keys) > 0:
self.send_error_report( await self.send_error_report(
metadata=metadata, metadata=metadata,
message=f'Failed to insert slots: \n {", ".join(error_keys)}', message=f'Failed to insert slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS', notification_id='REPORT_ORCHESTRATION_INSERTED_SLOTS',
@@ -584,14 +602,14 @@ class Formatters(BaseActivity):
success_keys, error_keys = self.parse_report(deleted_slots) success_keys, error_keys = self.parse_report(deleted_slots)
if len(success_keys) > 0: if len(success_keys) > 0:
self.send_success_report( await self.send_success_report(
metadata=metadata, metadata=metadata,
message=f'Deleted slots: \n {", ".join(success_keys)}', message=f'Deleted slots: \n {", ".join(success_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS', notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',
) )
if len(error_keys) > 0: if len(error_keys) > 0:
self.send_error_report( await self.send_error_report(
metadata=metadata, metadata=metadata,
message=f'Failed to delete slots: \n {", ".join(error_keys)}', message=f'Failed to delete slots: \n {", ".join(error_keys)}',
notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS', notification_id='REPORT_ORCHESTRATION_DELETED_SLOTS',

View File

@@ -6,41 +6,14 @@ with workflow.unsafe.imports_passed_through():
from logging import Logger from logging import Logger
from typing import Any from typing import Any
from pymongo import MongoClient
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.repository.mongodb_repository import MongoDBRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
def clear_mongo_id(docs: list) -> list: class MongoDB(SientiaMonitoring):
"""
Remove MongoDB internal `_id` fields from nested structures.
Args:
docs (list): The list of documents or nested structures to clean.
Returns:
list: The cleaned documents with `_id` fields removed wherever present.
"""
for doc in docs:
if isinstance(doc, list):
clear_mongo_id(doc)
elif isinstance(doc, dict):
if '_id' in doc:
del doc['_id']
for _key, value in doc.items():
if isinstance(value, list):
clear_mongo_id(value)
elif isinstance(value, dict):
clear_mongo_id([value])
return docs
class MongoDB(BaseActivity):
""" """
MongoDB operations activity for Temporal workflows. MongoDB operations activity for Temporal workflows.
@@ -64,60 +37,43 @@ class MongoDB(BaseActivity):
ttl_index_seconds: int, ttl_index_seconds: int,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, 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.connection_string = connection_string
self.database_name = database_name self.database_name = database_name
self.client: MongoClient = MongoClient( self.mongo_db_repository = MongoDBRepository(
self.connection_string, serverSelectionTimeoutMS=5000 connection_string=connection_string,
database_name=database_name,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
self.client.server_info() # Force early failure if connection is invalid
self.database = self.client[self.database_name]
self.ttl_index_seconds = ttl_index_seconds self.ttl_index_seconds = ttl_index_seconds
# Initialize MongoDB client here (omitted for brevity) # Initialize MongoDB client here (omitted for brevity)
logger.info('MongoDB connection initialized') logger.info('MongoDB connection initialized')
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) def close(self):
def shutdown(self):
""" """
Shutdown the MongoDB client and clean up resources. Shutdown the MongoDB client and clean up resources.
""" """
try: self.mongo_db_repository.close()
if self.client: SientiaMonitoring.shutdown(self)
self.logger.info('Closing MongoDB connection...')
self.client.close()
self.logger.info('MongoDB connection closed successfully')
except Exception as e:
self.logger.error(f'Failed to close MongoDB connection: {e}')
def __del__(self): def __del__(self):
""" """
Ensure the MongoDB client is closed when the object is garbage-collected. Ensure the MongoDB client is closed when the object is garbage-collected.
""" """
self.shutdown() self.close()
def find(self, collection_name: str, filters: dict[str, Any]) -> list[dict[str, Any]]:
"""
Find documents in a MongoDB collection based on the provided filters.
Args:
collection_name (str): The name of the collection to search in.
filters (dict[str, Any]): The query filters to apply.
Returns:
list[dict[str, Any]]: Documents matching the filters (with `_id` removed).
"""
collection = self.database[collection_name]
documents = list(collection.find(filters, {'_id': 0}))
documents = clear_mongo_id(documents)
return documents
@activity.defn( @activity.defn(
name='find_documents_in_mongodb', name='find_documents_in_mongodb',
@@ -151,7 +107,7 @@ class MongoDB(BaseActivity):
) )
try: try:
documents = self.find(collection_name, filters) documents = await self.mongo_db_repository.find(collection_name, filters, metadata)
self.info( self.info(
f"Loaded {len(documents)} documents from collection '{collection_name}'", f"Loaded {len(documents)} documents from collection '{collection_name}'",
@@ -173,7 +129,7 @@ class MongoDB(BaseActivity):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_QUERY_ERROR', notification_id='MONGODB_QUERY_ERROR',
message=f'Failed to execute MongoDB query: {e}', message=f'Failed to execute MongoDB query: {e}',
@@ -219,11 +175,9 @@ class MongoDB(BaseActivity):
) )
try: try:
collection = self.database[collection_name] aggregated_documents = await self.mongo_db_repository.aggregate(
collection_name, aggregation, metadata
aggregated_documents = list(collection.aggregate(aggregation)) )
aggregated_documents = clear_mongo_id(aggregated_documents)
self.info( self.info(
f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'", f"Aggregated {len(aggregated_documents)} documents from collection '{collection_name}'",
@@ -245,7 +199,7 @@ class MongoDB(BaseActivity):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_AGGREGATION_ERROR', notification_id='MONGODB_AGGREGATION_ERROR',
message=f'Failed to execute MongoDB aggregation: {e}', message=f'Failed to execute MongoDB aggregation: {e}',
@@ -268,12 +222,9 @@ class MongoDB(BaseActivity):
updated_pipelines = input_data.get('updated_pipelines', []) updated_pipelines = input_data.get('updated_pipelines', [])
metadata = input_data.get('metadata', {}) metadata = input_data.get('metadata', {})
date_now = now() date_now = now()
collection = self.database['orchestrated_schedules']
self.info('Updating pipelines timestamps...', metadata=metadata) self.info('Updating pipelines timestamps...', metadata=metadata)
success_count = 0
argument = [ argument = [
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']} {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
for pipeline in updated_pipelines for pipeline in updated_pipelines
@@ -282,11 +233,12 @@ class MongoDB(BaseActivity):
data_filter = {'$or': argument} if argument else {} data_filter = {'$or': argument} if argument else {}
try: try:
collection.update_many(data_filter, {'$set': {'updated_at': date_now}}) await self.mongo_db_repository.update_many(
success_count += 1 'orchestrated_schedules', data_filter, {'$set': {'updated_at': date_now}}, metadata
)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_UPDATE_PIPELINES_ERROR', notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message=f'Failed to update pipelines timestamps: {e}', message=f'Failed to update pipelines timestamps: {e}',
@@ -298,7 +250,7 @@ class MongoDB(BaseActivity):
raise e raise e
self.info( self.info(
f'Updated {success_count} of {len(updated_pipelines)} pipelines timestamps', f'Updated {len(updated_pipelines)} pipelines timestamps',
metadata=metadata, metadata=metadata,
) )
@@ -312,12 +264,9 @@ class MongoDB(BaseActivity):
""" """
created_pipelines = input_data.get('created_pipelines', []) created_pipelines = input_data.get('created_pipelines', [])
metadata = input_data.get('metadata', {}) metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info('Creating pipelines timestamps...', metadata=metadata) self.info('Creating pipelines timestamps...', metadata=metadata)
success_count = 0
date_now = now() date_now = now()
argument = [ argument = [
@@ -333,11 +282,12 @@ class MongoDB(BaseActivity):
try: try:
if data_filter: if data_filter:
collection.insert_many(data_filter) await self.mongo_db_repository.insert_many(
success_count += 1 'orchestrated_schedules', data_filter, metadata
)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_CREATE_PIPELINES_ERROR', notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message=f'Failed to create pipelines timestamps: {e}', message=f'Failed to create pipelines timestamps: {e}',
@@ -349,7 +299,7 @@ class MongoDB(BaseActivity):
raise e raise e
self.info( self.info(
f'Created {success_count} of {len(created_pipelines)} pipelines timestamps', f'Created {len(created_pipelines)} pipelines timestamps',
metadata=metadata, metadata=metadata,
) )
@@ -363,12 +313,9 @@ class MongoDB(BaseActivity):
""" """
deleted_pipelines = input_data.get('deleted_pipelines', []) deleted_pipelines = input_data.get('deleted_pipelines', [])
metadata = input_data.get('metadata', {}) metadata = input_data.get('metadata', {})
collection = self.database['orchestrated_schedules']
self.info('Deleting pipelines timestamps...', metadata=metadata) self.info('Deleting pipelines timestamps...', metadata=metadata)
success_count = 0
argument = [ argument = [
{'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']} {'schedule_name': pipeline['schedule_name'], 'namespace': pipeline['namespace']}
for pipeline in deleted_pipelines for pipeline in deleted_pipelines
@@ -377,11 +324,12 @@ class MongoDB(BaseActivity):
data_filter = {'$or': argument} if argument else {} data_filter = {'$or': argument} if argument else {}
try: try:
collection.delete_many(data_filter) await self.mongo_db_repository.delete_many(
success_count += 1 'orchestrated_schedules', data_filter, metadata
)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_DELETE_PIPELINES_ERROR', notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message=f'Failed to delete pipelines timestamps: {e}', message=f'Failed to delete pipelines timestamps: {e}',
@@ -393,7 +341,7 @@ class MongoDB(BaseActivity):
raise e raise e
self.info( self.info(
f'Deleted {success_count} of {len(deleted_pipelines)} pipelines timestamps', f'Deleted {len(deleted_pipelines)} pipelines timestamps',
metadata=metadata, metadata=metadata,
) )
@@ -413,7 +361,7 @@ class MongoDB(BaseActivity):
metadata=metadata, metadata=metadata,
) )
collection_names = self.database.list_collection_names() collection_names = self.mongo_db_repository.database.list_collection_names()
created_collections = [] created_collections = []
created_indexes = [] created_indexes = []
@@ -424,10 +372,10 @@ class MongoDB(BaseActivity):
try: try:
# Check if collection exists # Check if collection exists
if collection not in collection_names: if collection not in collection_names:
self.database.create_collection(collection) self.mongo_db_repository.database.create_collection(collection)
created_collections.append(collection) created_collections.append(collection)
collection = self.database[collection] collection = self.mongo_db_repository.database[collection]
# Check if TTL index exists # Check if TTL index exists
existing_indexes = collection.list_indexes() existing_indexes = collection.list_indexes()
ttl_index_exists = False ttl_index_exists = False
@@ -448,7 +396,7 @@ class MongoDB(BaseActivity):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGODB_CREATE_COLLECTION_ERROR', notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message=f'Failed to create collection {collection} with TTL index: {e}', message=f'Failed to create collection {collection} with TTL index: {e}',
@@ -510,7 +458,7 @@ class MongoDB(BaseActivity):
self.debug(f'Data filter: {data_filter}', metadata=metadata) self.debug(f'Data filter: {data_filter}', metadata=metadata)
data = self.find(collection_name, data_filter) data = await self.mongo_db_repository.find(collection_name, data_filter, metadata)
self.debug(f'Collected: {data}', metadata=metadata) self.debug(f'Collected: {data}', metadata=metadata)
@@ -526,7 +474,7 @@ class MongoDB(BaseActivity):
return data return data
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='MONGO_LOAD_ERROR', notification_id='MONGO_LOAD_ERROR',
message=f'Error loading data from MongoDB: {e}', message=f'Error loading data from MongoDB: {e}',

View File

@@ -10,11 +10,12 @@ with workflow.unsafe.imports_passed_through():
from pandas import DataFrame from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.redis_base import Redis from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from sientia_do.repository.redis_repository import RedisRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ, now
class SlotManager(Redis): class SlotManager(SientiaMonitoring):
""" """
Redis-based OPC slot management and notification filtering activity. Redis-based OPC slot management and notification filtering activity.
@@ -48,8 +49,37 @@ class SlotManager(Redis):
password: str, password: str,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController,
): ):
Redis.__init__(self, host, port, username, password, logger, notification_handler) 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.
"""
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') @activity.defn(name='load_opc_slots')
async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]: async def load_opc_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -80,7 +110,7 @@ class SlotManager(Redis):
opc_slots = {} opc_slots = {}
try: try:
slot_keys = self.redis_client.keys('slot:opc_tags:*') slot_keys = await self.redis_repository.keys('slot:opc_tags:*')
self.debug(f'Slot keys: {slot_keys}', metadata=metadata) self.debug(f'Slot keys: {slot_keys}', metadata=metadata)
@@ -91,10 +121,10 @@ class SlotManager(Redis):
decoded_keys = slot_keys decoded_keys = slot_keys
for key in decoded_keys: for key in decoded_keys:
opc_slots[key] = self.get(key) opc_slots[key] = await self.redis_repository.get(key)
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message=f'Failed to load OPC slots: {e}', message=f'Failed to load OPC slots: {e}',
@@ -123,7 +153,7 @@ class SlotManager(Redis):
self.info('Loading active ingestors...', metadata=metadata) self.info('Loading active ingestors...', metadata=metadata)
try: try:
active_ingestors = self.redis_client.keys('heartbeat:ingestor:*') active_ingestors = await self.redis_repository.keys('heartbeat:ingestor:*')
self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata) self.info(f'Loaded {len(active_ingestors)} active ingestors', metadata=metadata)
@@ -140,7 +170,7 @@ class SlotManager(Redis):
return ingestors return ingestors
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message=f'Failed to load active ingestors: {e}', message=f'Failed to load active ingestors: {e}',
@@ -175,7 +205,7 @@ class SlotManager(Redis):
for slot in to_insert: for slot in to_insert:
try: try:
self.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None) await self.redis_repository.set(f'slot:opc_tags:{slot}', to_insert[slot], ttl=None)
report[slot] = {'success': True, 'message': 'Slot updated successfully'} report[slot] = {'success': True, 'message': 'Slot updated successfully'}
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
@@ -213,7 +243,7 @@ class SlotManager(Redis):
for slot in to_delete: for slot in to_delete:
try: try:
self.redis_client.delete(f'slot:opc_tags:{slot}') await self.redis_repository.delete(f'slot:opc_tags:{slot}')
report[slot] = {'success': True, 'message': 'Slot deleted successfully'} report[slot] = {'success': True, 'message': 'Slot deleted successfully'}
success_count += 1 success_count += 1
except Exception as e: except Exception as e:
@@ -243,9 +273,9 @@ class SlotManager(Redis):
key = f'notification_last_timestamp:{input_data["mail_type"]}' key = f'notification_last_timestamp:{input_data["mail_type"]}'
try: try:
data_hold = self.get(key) data_hold = await self.redis_repository.get(key)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message=f'Error getting last data timestamp: {e}', message=f'Error getting last data timestamp: {e}',
@@ -290,9 +320,9 @@ class SlotManager(Redis):
self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata) self.debug(f'Last collected timestamp to insert: {last_data_timestamp}', metadata=metadata)
try: try:
self.set(key, last_data_timestamp, ttl=60 * 60 * 5) await self.redis_repository.set(key, last_data_timestamp, ttl=60 * 60 * 5)
except Exception as e: except Exception as e:
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message=f'Error setting last data timestamp: {e}', message=f'Error setting last data timestamp: {e}',
@@ -361,7 +391,7 @@ class SlotManager(Redis):
# Check if notification was recently sent # Check if notification was recently sent
key = f'{notification["trigger"]}:{notification_id}' key = f'{notification["trigger"]}:{notification_id}'
last_sent = self.get(key) last_sent = await self.redis_repository.get(key)
if last_sent is None: if last_sent is None:
alert_type = 'core_alerts' alert_type = 'core_alerts'
@@ -414,6 +444,6 @@ class SlotManager(Redis):
status = row['status'] status = row['status']
if status == 'sent': if status == 'sent':
key = f'{row["schedule"]}:{row["notification_id"]}' key = f'{row["schedule"]}:{row["notification_id"]}'
self.set(key, date_now, ttl=sent_ttl) await self.redis_repository.set(key, date_now, ttl=sent_ttl)
self.info('Notification cache stored...', metadata=metadata) self.info('Notification cache stored...', metadata=metadata)

View File

@@ -19,12 +19,12 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.base import BaseActivity from sientia_do.observability.sientia_monitoring import MetricsController, SientiaMonitoring
from orchestrator.utils.converters import parse_frequency from orchestrator.utils.converters import parse_frequency
class TemporalManager(BaseActivity): class TemporalManager(SientiaMonitoring):
""" """
Temporal workflow and schedule management activity. Temporal workflow and schedule management activity.
@@ -48,7 +48,14 @@ class TemporalManager(BaseActivity):
laborious_namespace: str, laborious_namespace: str,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController,
): ):
SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
self.temporal_host = host self.temporal_host = host
self.scouter_namespace = scouter_namespace self.scouter_namespace = scouter_namespace
self.laborious_namespace = laborious_namespace self.laborious_namespace = laborious_namespace
@@ -58,7 +65,17 @@ class TemporalManager(BaseActivity):
self.model_name_id_key = SearchAttributeKey.for_keyword('model_name') self.model_name_id_key = SearchAttributeKey.for_keyword('model_name')
self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated') self.orchestrated_id_key = SearchAttributeKey.for_keyword('orchestrated')
BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) def close(self):
"""
Close the TemporalManager connection and clean up 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): async def connect_to_temporal(self):
""" """
@@ -127,7 +144,7 @@ class TemporalManager(BaseActivity):
except Exception as e: except Exception as e:
trace = traceback.format_exc() trace = traceback.format_exc()
self.send_notification( await self.send_notification_async(
metadata=metadata, metadata=metadata,
notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR', notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR',
message=f'Failed to normalize schedules: {e}', message=f'Failed to normalize schedules: {e}',

View File

@@ -114,11 +114,7 @@ python_classes = ["Test*"]
python_functions = ["test_*"] python_functions = ["test_*"]
addopts = [ addopts = [
"-v", "-v",
"--strict-markers", "--strict-markers"
"--cov=model_manager",
"--cov-report=term-missing",
"--cov-report=html",
"--cov-report=xml",
] ]
markers = [ markers = [
"asyncio: marks tests as async", "asyncio: marks tests as async",

View File

@@ -2,8 +2,7 @@ temporalio
psycopg2-binary psycopg2-binary
sqlalchemy sqlalchemy
redis redis
couchbase
pymongo pymongo
jinja2 jinja2
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.4.7 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.5.3
prometheus-client prometheus-client

View File

@@ -13,7 +13,9 @@ from orchestrator.activities.temporal_manager import TemporalManager
@patch('orchestrator.activities.formatters.Formatters.__init__') @patch('orchestrator.activities.formatters.Formatters.__init__')
@patch('orchestrator.activities.email.Email.__init__') @patch('orchestrator.activities.email.Email.__init__')
@patch('sientia_do.temporal.activities.postgres.Postgres.__init__') @patch('sientia_do.temporal.activities.postgres.Postgres.__init__')
@patch('orchestrator.activities.activities.MetricsController')
def test___init__( def test___init__(
mock_metrics_controller,
mock_postgres_init, mock_postgres_init,
mock_email_init, mock_email_init,
mock_formatters_init, mock_formatters_init,
@@ -79,6 +81,7 @@ def test___init__(
password='password', password='password',
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
) )
mock_mongodb_init.assert_called_once_with( mock_mongodb_init.assert_called_once_with(
@@ -88,6 +91,7 @@ def test___init__(
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
) )
mock_temporal_manager_init.assert_called_once_with( mock_temporal_manager_init.assert_called_once_with(
@@ -97,6 +101,7 @@ def test___init__(
laborious_namespace='laborious', laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
) )
mock_formatters_init.assert_called_once_with( mock_formatters_init.assert_called_once_with(
@@ -105,28 +110,23 @@ def test___init__(
laborious_namespace='laborious', laborious_namespace='laborious',
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=mock_metrics_controller.return_value,
) )
@patch('orchestrator.activities.mongo_db.MongoDB.__init__') @patch('orchestrator.activities.activities.MongoDB')
@patch('orchestrator.activities.temporal_manager.TemporalManager.__init__') @patch('orchestrator.activities.activities.TemporalManager')
@patch('orchestrator.activities.slot_manager.SlotManager.__init__') @patch('orchestrator.activities.activities.SlotManager')
@patch('orchestrator.activities.formatters.Formatters.__init__') @patch('orchestrator.activities.activities.Formatters')
@patch('orchestrator.activities.email.Email.__init__') @patch('orchestrator.activities.activities.Email')
@patch('sientia_do.temporal.activities.postgres.Postgres.__init__') @patch('orchestrator.activities.activities.Postgres')
@patch('orchestrator.activities.email.Email.shutdown')
@patch('sientia_do.temporal.activities.postgres.Postgres.close')
@patch('orchestrator.activities.mongo_db.MongoDB.shutdown')
def test_shutdown( def test_shutdown(
mock_mongodb_close, mock_mongodb,
mock_postgres_shutdown, mock_temporal_manager,
mock_email_close, mock_slot_manager,
mock_postgres_init, mock_formatters,
mock_email_init, mock_email,
mock_formatters_init, mock_postgres,
mock_slot_manager_init,
mock_temporal_manager_init,
mock_mongodb_init,
): ):
activities = Activities( activities = Activities(
temporal_config=MagicMock(), temporal_config=MagicMock(),
@@ -140,6 +140,9 @@ def test_shutdown(
activities.shutdown() activities.shutdown()
mock_mongodb_close.assert_called() mock_mongodb.close.assert_called()
mock_postgres_shutdown.assert_called() mock_temporal_manager.close.assert_called()
mock_email_close.assert_called() mock_slot_manager.close.assert_called()
mock_formatters.close.assert_called()
mock_email.close.assert_called()
mock_postgres.close.assert_called()

View File

@@ -1,75 +0,0 @@
from unittest.mock import ANY, MagicMock, patch
from pytest import fixture, mark, raises
from sientia_do.notifications.models import NotificationLevel
from orchestrator.activities.couchbase import Couchbase
@fixture
@patch('orchestrator.activities.couchbase.Cluster')
def couchbase(_cluster_mock):
return Couchbase(
connection_string='couchbase://localhost',
username='admin',
password='password',
logger=MagicMock(),
notification_handler=MagicMock(),
)
def test_shutdown_success(couchbase):
couchbase.shutdown()
couchbase.cluster.close.assert_called_once()
def test_shutdown_failure(couchbase):
couchbase.cluster.close.side_effect = Exception('Test error')
couchbase.shutdown()
couchbase.logger.error.assert_called_once_with(
f'Failed to close Couchbase connection: {couchbase.cluster.close.side_effect}'
)
@mark.asyncio
async def test_load_query_from_couchbase_success(couchbase):
couchbase.cluster.query.return_value.rows.return_value = [
{'id': '1', 'name': 'test'},
{'id': '2', 'name': 'test2'},
]
query = 'SELECT * FROM bucket'
result = await couchbase.load_query_from_couchbase(
{
'query': query,
}
)
assert result == [
{'id': '1', 'name': 'test'},
{'id': '2', 'name': 'test2'},
]
couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_not_called()
@mark.asyncio
async def test_load_query_from_couchbase_failure(couchbase):
couchbase.cluster.query.side_effect = ValueError('Test error')
query = 'SELECT * FROM bucket'
with raises(ValueError):
await couchbase.load_query_from_couchbase(
{
'query': query,
}
)
couchbase.cluster.query.assert_called_once_with(query)
couchbase.notification_handler.build_and_send_notification.assert_called_once_with(
notification_id='COUCHBASE_LOAD_QUERY_ERROR',
message='Failed to execute couchbase query: Test error',
block='load_query_from_couchbase',
level=NotificationLevel.ERROR,
attachment_content=ANY,
)

View File

@@ -1,5 +1,5 @@
from smtplib import SMTPServerDisconnected from smtplib import SMTPServerDisconnected
from unittest.mock import MagicMock, call, patch from unittest.mock import AsyncMock, MagicMock, call, patch
from pytest import fixture, mark from pytest import fixture, mark
@@ -17,8 +17,13 @@ def email(smtplib, email_builder):
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
email_builder.send_notification_async = AsyncMock()
email_builder.send_notification = MagicMock() email_builder.send_notification = MagicMock()
email.send_notification_async = AsyncMock()
email.send_notification = MagicMock()
email.emit_metric = AsyncMock()
return email return email
@@ -33,6 +38,7 @@ def test___init___with_password(smtplib, email_builder):
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
assert email.sender_email == 'test@test.com' assert email.sender_email == 'test@test.com'
@@ -56,6 +62,7 @@ def test___init___without_password(smtplib, email_builder):
smtp_port=587, smtp_port=587,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
assert email.sender_email == 'test@test.com' assert email.sender_email == 'test@test.com'
@@ -76,10 +83,12 @@ metadata = {
} }
def test_shutdown(email): @patch('orchestrator.activities.email.SientiaMonitoring')
email.shutdown() def test_close(sientia_monitoring_mock, email):
email.close()
email.server.quit.assert_called_once() email.server.quit.assert_called_once()
sientia_monitoring_mock.shutdown.assert_called_once()
@mark.asyncio @mark.asyncio

View File

@@ -1,5 +1,5 @@
import json import json
from unittest.mock import MagicMock, call, patch from unittest.mock import AsyncMock, MagicMock, call, patch
from pandas import DataFrame from pandas import DataFrame
from pytest import fixture, mark from pytest import fixture, mark
@@ -15,9 +15,12 @@ def formatters():
laborious_namespace='laborious', laborious_namespace='laborious',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
formatters.send_notification = MagicMock() formatters.send_notification = MagicMock()
formatters.send_notification_async = AsyncMock()
formatters.emit_metric = AsyncMock()
return formatters return formatters
@@ -173,7 +176,7 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma
] ]
) )
formatters.send_notification.assert_has_calls( formatters.send_notification_async.assert_has_calls(
[ [
call( call(
metadata=metadata['metadata'], metadata=metadata['metadata'],
@@ -288,14 +291,15 @@ async def test_create_slot_config(formatters):
} }
def test_send_success_report(formatters): @mark.asyncio
formatters.send_success_report( async def test_send_success_report(formatters):
await formatters.send_success_report(
metadata=metadata, metadata=metadata,
message='test_message', message='test_message',
notification_id='test_notification_id', notification_id='test_notification_id',
attachment={'test': 'test'}, attachment={'test': 'test'},
) )
formatters.send_notification.assert_called_once_with( formatters.send_notification_async.assert_called_once_with(
metadata=metadata, metadata=metadata,
notification_id='test_notification_id', notification_id='test_notification_id',
message='test_message', message='test_message',
@@ -305,14 +309,15 @@ def test_send_success_report(formatters):
) )
def test_send_error_report(formatters): @mark.asyncio
formatters.send_error_report( async def test_send_error_report(formatters):
await formatters.send_error_report(
metadata=metadata, metadata=metadata,
message='test_message', message='test_message',
notification_id='test_notification_id', notification_id='test_notification_id',
attachment='test_attachment', attachment='test_attachment',
) )
formatters.send_notification.assert_called_once_with( formatters.send_notification_async.assert_called_once_with(
metadata=metadata, metadata=metadata,
notification_id='test_notification_id', notification_id='test_notification_id',
message='test_message', message='test_message',
@@ -364,8 +369,8 @@ def test_parse_report_schedule(formatters):
@mark.asyncio @mark.asyncio
async def test_report_schedule_orchestration(formatters): async def test_report_schedule_orchestration(formatters):
formatters.parse_report_schedule = MagicMock(side_effect=formatters.parse_report_schedule) formatters.parse_report_schedule = MagicMock(side_effect=formatters.parse_report_schedule)
formatters.send_success_report = MagicMock() formatters.send_success_report = AsyncMock()
formatters.send_error_report = MagicMock() formatters.send_error_report = AsyncMock()
input_data = { input_data = {
**metadata, **metadata,
@@ -496,8 +501,8 @@ async def test_report_schedule_orchestration(formatters):
@mark.asyncio @mark.asyncio
async def test_report_slot_orchestration(formatters): async def test_report_slot_orchestration(formatters):
formatters.parse_report = MagicMock(side_effect=formatters.parse_report) formatters.parse_report = MagicMock(side_effect=formatters.parse_report)
formatters.send_success_report = MagicMock() formatters.send_success_report = AsyncMock()
formatters.send_error_report = MagicMock() formatters.send_error_report = AsyncMock()
input_data = { input_data = {
**metadata, **metadata,

View File

@@ -1,47 +1,15 @@
from datetime import datetime from datetime import datetime
from unittest.mock import ANY, MagicMock, patch from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import fixture, mark from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ from sientia_do.temporal.constants import DATETIME_FORMAT_MS_WITH_TZ
from orchestrator.activities.mongo_db import MongoDB, clear_mongo_id from orchestrator.activities.mongo_db import MongoDB
def test_clear_mongo_id():
input_data = [
[
{
'name': 'test',
'_id': '12345',
}
],
{
'name': 'test',
'_id': '12345',
'nested': {
'_id': '67890',
'value': [1, 2, 3],
'list': [{'_id': 'abcde', 'item': 'value'}],
},
'nested_list': [{'_id': 'fghij', 'item': 'value1'}, {'_id': 'klmno', 'item': 'value2'}],
},
]
output = clear_mongo_id(input_data)
assert output == [
[{'name': 'test'}],
{
'name': 'test',
'nested': {'value': [1, 2, 3], 'list': [{'item': 'value'}]},
'nested_list': [{'item': 'value1'}, {'item': 'value2'}],
},
]
@fixture @fixture
@patch('orchestrator.activities.mongo_db.MongoClient') @patch('orchestrator.activities.mongo_db.MongoDBRepository')
def mongo_db(mongo_mock): def mongo_db(mongo_mock):
mongo = MongoDB( mongo = MongoDB(
connection_string='mongodb://localhost:27017', connection_string='mongodb://localhost:27017',
@@ -49,63 +17,72 @@ def mongo_db(mongo_mock):
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
mongo.send_notification = MagicMock() mongo.send_notification = MagicMock()
mongo.send_notification_async = AsyncMock()
mongo.emit_metric = AsyncMock()
return mongo return mongo
@patch('orchestrator.activities.mongo_db.MongoClient') @patch('orchestrator.activities.mongo_db.MongoDBRepository')
def test___init__(mongo_mock): def test___init__(mongo_mock):
mongo_db = MongoDB( logger = MagicMock()
notification_handler = MagicMock()
metrics_controller = AsyncMock()
MongoDB(
connection_string='mongodb://localhost:27017', connection_string='mongodb://localhost:27017',
database_name='test_db', database_name='test_db',
ttl_index_seconds=3600, ttl_index_seconds=3600,
logger=MagicMock(), logger=logger,
notification_handler=MagicMock(), notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
assert mongo_db.connection_string == 'mongodb://localhost:27017'
assert mongo_db.database_name == 'test_db'
mongo_mock.assert_called_once_with('mongodb://localhost:27017', serverSelectionTimeoutMS=5000)
mongo_db.client.server_info.assert_called_once()
mongo_db.client.__getitem__.assert_called_once_with('test_db')
mongo_mock.assert_called_once_with(
def test_shutdown_success(mongo_db): connection_string='mongodb://localhost:27017',
mongo_db.shutdown() database_name='test_db',
mongo_db.client.close.assert_called_once() logger=logger,
mongo_db.logger.info.assert_any_call('Closing MongoDB connection...') notification_handler=notification_handler,
mongo_db.logger.info.assert_any_call('MongoDB connection closed successfully') metrics_controller=metrics_controller,
def test_shutdown_failure(mongo_db):
mongo_db.client.close.side_effect = Exception('Close failed')
mongo_db.shutdown()
mongo_db.logger.error.assert_called_once_with(
'Failed to close MongoDB connection: Close failed'
) )
@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()
@mark.asyncio @mark.asyncio
async def test_find_documents_in_mongodb_success(mongo_db): async def test_find_documents_in_mongodb_success(mongo_db):
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}} input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
mock_collection = MagicMock()
mock_collection.find.return_value = [ mongo_db.mongo_db_repository.find = AsyncMock(
{ return_value=[
'_id': '12345', {
'name': 'test1', 'name': 'test1',
'timestamp': datetime.strptime( 'timestamp': datetime.strptime(
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
), ),
}, },
{ {
'_id': '67890', 'name': 'test2',
'name': 'test2', 'timestamp': datetime.strptime(
'timestamp': datetime.strptime( '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ ),
), },
}, ]
] )
mongo_db.database.__getitem__.return_value = mock_collection
result = await mongo_db.find_documents_in_mongodb( result = await mongo_db.find_documents_in_mongodb(
{'query': input_data, 'timestamp_fields': ['timestamp']} {'query': input_data, 'timestamp_fields': ['timestamp']}
@@ -114,7 +91,9 @@ async def test_find_documents_in_mongodb_success(mongo_db):
assert len(result) == 2 assert len(result) == 2
assert result[0] == {'name': 'test1', 'timestamp': '2023-01-01 12:00:00.000000+0000'} 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'} assert result[1] == {'name': 'test2', 'timestamp': '2023-01-01 12:00:00.000000+0000'}
mock_collection.find.assert_called_once_with({'name': {'$exists': True}}, {'_id': 0}) mongo_db.mongo_db_repository.find.assert_called_once_with(
'test_collection', {'name': {'$exists': True}}, {}
)
metadata = { metadata = {
@@ -130,15 +109,13 @@ metadata = {
@mark.asyncio @mark.asyncio
async def test_find_documents_in_mongodb_failure(mongo_db): async def test_find_documents_in_mongodb_failure(mongo_db):
input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}} input_data = {'collection': 'test_collection', 'filters': {'name': {'$exists': True}}}
mongo_db.database.__getitem__.return_value = MagicMock( mongo_db.mongo_db_repository.find = AsyncMock(side_effect=Exception('Error'))
find=MagicMock(side_effect=Exception('Error'))
)
try: try:
await mongo_db.find_documents_in_mongodb({'query': input_data, **metadata}) await mongo_db.find_documents_in_mongodb({'query': input_data, **metadata})
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_QUERY_ERROR', notification_id='MONGODB_QUERY_ERROR',
message='Failed to execute MongoDB query: Error', message='Failed to execute MongoDB query: Error',
@@ -171,24 +148,22 @@ async def test_aggregate_documents_in_mongodb_success(mongo_db):
'collection': 'test_collection', 'collection': 'test_collection',
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}], 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
} }
mock_collection = MagicMock() mongo_db.mongo_db_repository.aggregate = AsyncMock(
mock_collection.aggregate.return_value = [ return_value=[
{ {
'_id': 'asdad', 'name': 'test1',
'name': 'test1', 'timestamp': datetime.strptime(
'timestamp': datetime.strptime( '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ ),
), },
}, {
{ 'name': 'test2',
'_id': 'adzx', 'timestamp': datetime.strptime(
'name': 'test2', '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
'timestamp': datetime.strptime( ),
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ },
), ]
}, )
]
mongo_db.database.__getitem__.return_value = mock_collection
result = await mongo_db.aggregate_documents_in_mongodb( result = await mongo_db.aggregate_documents_in_mongodb(
{'query': input_data, 'timestamp_fields': ['timestamp']} {'query': input_data, 'timestamp_fields': ['timestamp']}
@@ -200,7 +175,9 @@ async def test_aggregate_documents_in_mongodb_success(mongo_db):
expected_pipeline = input_data['aggregation'] expected_pipeline = input_data['aggregation']
expected_pipeline.append({'$project': {'_id': 0}}) expected_pipeline.append({'$project': {'_id': 0}})
mock_collection.aggregate.assert_called_once_with(expected_pipeline) mongo_db.mongo_db_repository.aggregate.assert_called_once_with(
'test_collection', expected_pipeline, {}
)
@mark.asyncio @mark.asyncio
@@ -209,15 +186,13 @@ async def test_aggregate_documents_in_mongodb_failure(mongo_db):
'collection': 'test_collection', 'collection': 'test_collection',
'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}], 'aggregation': [{'$match': {'name': {'$exists': True}}}, {'$project': {'name': 1}}],
} }
mongo_db.database.__getitem__.return_value = MagicMock( mongo_db.mongo_db_repository.aggregate = AsyncMock(side_effect=Exception('Error'))
aggregate=MagicMock(side_effect=Exception('Error'))
)
try: try:
await mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata}) await mongo_db.aggregate_documents_in_mongodb({'query': input_data, **metadata})
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_AGGREGATION_ERROR', notification_id='MONGODB_AGGREGATION_ERROR',
message='Failed to execute MongoDB aggregation: Error', message='Failed to execute MongoDB aggregation: Error',
@@ -267,9 +242,10 @@ async def test_update_pipelines_timestamps_success(now_mock, mongo_db):
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
] ]
} }
mongo_db.database['pipelines'].update_many.return_value = MagicMock() mongo_db.mongo_db_repository.update_many = AsyncMock(return_value=MagicMock())
await mongo_db.update_pipelines_timestamps(input_data) await mongo_db.update_pipelines_timestamps(input_data)
mongo_db.database['pipelines'].update_many.assert_called_once_with( mongo_db.mongo_db_repository.update_many.assert_called_once_with(
'orchestrated_schedules',
{ {
'$or': [ '$or': [
{'schedule_name': 'test1', 'namespace': 'test1'}, {'schedule_name': 'test1', 'namespace': 'test1'},
@@ -277,6 +253,7 @@ async def test_update_pipelines_timestamps_success(now_mock, mongo_db):
] ]
}, },
{'$set': {'updated_at': now_mock.return_value}}, {'$set': {'updated_at': now_mock.return_value}},
{},
) )
@@ -291,12 +268,12 @@ async def test_update_pipelines_timestamps_failure(now_mock, mongo_db):
**metadata, **metadata,
} }
mongo_db.database['pipelines'].update_many.side_effect = Exception('Error') mongo_db.mongo_db_repository.update_many.side_effect = Exception('Error')
try: try:
await mongo_db.update_pipelines_timestamps(input_data) await mongo_db.update_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_UPDATE_PIPELINES_ERROR', notification_id='MONGODB_UPDATE_PIPELINES_ERROR',
message='Failed to update pipelines timestamps: Error', message='Failed to update pipelines timestamps: Error',
@@ -318,13 +295,15 @@ async def test_create_pipelines_timestamps_success(now_mock, mongo_db):
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
] ]
} }
mongo_db.database['pipelines'].insert_many.return_value = MagicMock() mongo_db.mongo_db_repository.insert_many = AsyncMock(return_value=MagicMock())
await mongo_db.create_pipelines_timestamps(input_data) await mongo_db.create_pipelines_timestamps(input_data)
mongo_db.database['pipelines'].insert_many.assert_called_once_with( 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': 'test1', 'namespace': 'test1', 'updated_at': now_mock.return_value},
{'schedule_name': 'test2', 'namespace': 'test2', 'updated_at': now_mock.return_value}, {'schedule_name': 'test2', 'namespace': 'test2', 'updated_at': now_mock.return_value},
] ],
{},
) )
@@ -338,12 +317,12 @@ async def test_create_pipelines_timestamps_failure(now_mock, mongo_db):
], ],
**metadata, **metadata,
} }
mongo_db.database['pipelines'].insert_many.side_effect = Exception('Error') mongo_db.mongo_db_repository.insert_many.side_effect = Exception('Error')
try: try:
await mongo_db.create_pipelines_timestamps(input_data) await mongo_db.create_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_CREATE_PIPELINES_ERROR', notification_id='MONGODB_CREATE_PIPELINES_ERROR',
message='Failed to create pipelines timestamps: Error', message='Failed to create pipelines timestamps: Error',
@@ -365,15 +344,17 @@ async def test_delete_pipelines_timestamps_success(now_mock, mongo_db):
{'schedule_name': 'test2', 'namespace': 'test2', 'success': True}, {'schedule_name': 'test2', 'namespace': 'test2', 'success': True},
] ]
} }
mongo_db.database['pipelines'].delete_many.return_value = MagicMock() mongo_db.mongo_db_repository.delete_many = AsyncMock(return_value=MagicMock())
await mongo_db.delete_pipelines_timestamps(input_data) await mongo_db.delete_pipelines_timestamps(input_data)
mongo_db.database['pipelines'].delete_many.assert_called_once_with( mongo_db.mongo_db_repository.delete_many.assert_called_once_with(
'orchestrated_schedules',
{ {
'$or': [ '$or': [
{'schedule_name': 'test1', 'namespace': 'test1'}, {'schedule_name': 'test1', 'namespace': 'test1'},
{'schedule_name': 'test2', 'namespace': 'test2'}, {'schedule_name': 'test2', 'namespace': 'test2'},
] ]
} },
{},
) )
@@ -387,12 +368,12 @@ async def test_delete_pipelines_timestamps_failure(now_mock, mongo_db):
], ],
**metadata, **metadata,
} }
mongo_db.database['pipelines'].delete_many.side_effect = Exception('Error') mongo_db.mongo_db_repository.delete_many.side_effect = Exception('Error')
try: try:
await mongo_db.delete_pipelines_timestamps(input_data) await mongo_db.delete_pipelines_timestamps(input_data)
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_DELETE_PIPELINES_ERROR', notification_id='MONGODB_DELETE_PIPELINES_ERROR',
message='Failed to delete pipelines timestamps: Error', message='Failed to delete pipelines timestamps: Error',
@@ -416,10 +397,12 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
}, },
} }
mongo_db.database.list_collection_names.return_value = [ mongo_db.mongo_db_repository.database.list_collection_names = MagicMock(
'raw_scouter_pipeline_2', return_value=[
'raw_scouter_pipeline_3', 'raw_scouter_pipeline_2',
] 'raw_scouter_pipeline_3',
]
)
collection_1 = MagicMock( collection_1 = MagicMock(
list_indexes=MagicMock( list_indexes=MagicMock(
@@ -438,15 +421,17 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': 3600}]) list_indexes=MagicMock(return_value=[{'key': 'inserted_at', 'expireAfterSeconds': 3600}])
) )
mongo_db.database.__getitem__ = MagicMock( mongo_db.mongo_db_repository.database.__getitem__ = MagicMock(
side_effect=[collection_1, collection_2, collection_3] side_effect=[collection_1, collection_2, collection_3]
) )
await mongo_db.create_collection_with_ttl_index(input_data) await mongo_db.create_collection_with_ttl_index(input_data)
mongo_db.database.list_collection_names.assert_called_once_with() mongo_db.mongo_db_repository.database.list_collection_names.assert_called_once_with()
mongo_db.database.create_collection.assert_called_once_with('raw_scouter_pipeline') mongo_db.mongo_db_repository.database.create_collection.assert_called_once_with(
'raw_scouter_pipeline'
)
collection_1.list_indexes.assert_called_once() collection_1.list_indexes.assert_called_once()
collection_1.create_index.assert_called_once_with( collection_1.create_index.assert_called_once_with(
@@ -466,15 +451,15 @@ async def test_create_collection_with_ttl_index_success(mongo_db):
async def test_create_collection_with_ttl_index_failure(mongo_db): async def test_create_collection_with_ttl_index_failure(mongo_db):
input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}} input_data = {**metadata, 'pipelines': {'scouter-pipeline': {'topic': 'raw_scouter_pipeline'}}}
mongo_db.database.list_collection_names.return_value = [] mongo_db.mongo_db_repository.database.list_collection_names.return_value = []
mongo_db.database.create_collection.side_effect = Exception('Error') mongo_db.mongo_db_repository.database.create_collection.side_effect = Exception('Error')
try: try:
await mongo_db.create_collection_with_ttl_index(input_data) await mongo_db.create_collection_with_ttl_index(input_data)
except Exception as e: except Exception as e:
assert str(e) == 'Error' assert str(e) == 'Error'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='MONGODB_CREATE_COLLECTION_ERROR', notification_id='MONGODB_CREATE_COLLECTION_ERROR',
message='Failed to create collection raw_scouter_pipeline with TTL index: Error', message='Failed to create collection raw_scouter_pipeline with TTL index: Error',
@@ -489,18 +474,17 @@ async def test_create_collection_with_ttl_index_failure(mongo_db):
@mark.asyncio @mark.asyncio
async def test_load_latest_data_none_last_data_timestamp(mongo_db): async def test_load_latest_data_none_last_data_timestamp(mongo_db):
"""Test load_latest_data""" """Test load_latest_data"""
collection = MagicMock() mongo_db.mongo_db_repository.find = AsyncMock(
mongo_db.database.__getitem__.return_value = collection return_value=[
{
collection.find.return_value = [ 'name': 'test1',
{ 'value': 1,
'name': 'test1', 'timestamp': datetime.strptime(
'value': 1, '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
'timestamp': datetime.strptime( ),
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ }
), ]
} )
]
result = await mongo_db.load_latest_data( result = await mongo_db.load_latest_data(
{ {
@@ -511,9 +495,11 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db):
} }
) )
mongo_db.database.__getitem__.assert_called_once_with('test_collection') mongo_db.mongo_db_repository.find.assert_called_once_with(
'test_collection',
collection.find.assert_called_once_with({'level': 'ERROR'}, {'_id': 0}) {'level': 'ERROR'},
{'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
)
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}] assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}]
@@ -521,18 +507,18 @@ async def test_load_latest_data_none_last_data_timestamp(mongo_db):
@mark.asyncio @mark.asyncio
async def test_load_latest_data_not_none_last_data_timestamp(mongo_db): async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
"""Test load_latest_data""" """Test load_latest_data"""
collection = MagicMock()
mongo_db.database.__getitem__.return_value = collection
collection.find.return_value = [ mongo_db.mongo_db_repository.find = AsyncMock(
{ return_value=[
'name': 'test1', {
'value': 1, 'name': 'test1',
'timestamp': datetime.strptime( 'value': 1,
'2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ 'timestamp': datetime.strptime(
), '2023-01-01 12:00:00.000000+0000', DATETIME_FORMAT_MS_WITH_TZ
} ),
] }
]
)
result = await mongo_db.load_latest_data( result = await mongo_db.load_latest_data(
{ {
@@ -543,9 +529,8 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
} }
) )
mongo_db.database.__getitem__.assert_called_once_with('test_collection') mongo_db.mongo_db_repository.find.assert_called_once_with(
'test_collection',
collection.find.assert_called_once_with(
{ {
'level': 'ERROR', 'level': 'ERROR',
'timestamp': { 'timestamp': {
@@ -554,7 +539,7 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
) )
}, },
}, },
{'_id': 0}, {'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
) )
assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}] assert result == [{'name': 'test1', 'value': 1, 'timestamp': '2023-01-01 12:00:00.000000+0000'}]
@@ -563,11 +548,7 @@ async def test_load_latest_data_not_none_last_data_timestamp(mongo_db):
@mark.asyncio @mark.asyncio
async def test_load_latest_data_error(mongo_db): async def test_load_latest_data_error(mongo_db):
"""Test load_latest_data""" """Test load_latest_data"""
collection = MagicMock() mongo_db.mongo_db_repository.find.side_effect = Exception('test')
mongo_db.send_notification = MagicMock()
mongo_db.database.__getitem__.return_value = collection
collection.find.side_effect = Exception('test')
try: try:
await mongo_db.load_latest_data( await mongo_db.load_latest_data(
@@ -581,7 +562,7 @@ async def test_load_latest_data_error(mongo_db):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
mongo_db.send_notification.assert_called_once_with( mongo_db.send_notification_async.assert_called_once_with(
metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'}, metadata={'workflow_name': 'test_pipeline', 'schedule_name': 'test_schedule'},
notification_id='MONGO_LOAD_ERROR', notification_id='MONGO_LOAD_ERROR',
message='Error loading data from MongoDB: test', message='Error loading data from MongoDB: test',

View File

@@ -1,5 +1,5 @@
from datetime import timedelta from datetime import timedelta
from unittest.mock import ANY, MagicMock, call, patch from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from pandas import DataFrame from pandas import DataFrame
from pytest import fixture, mark from pytest import fixture, mark
@@ -19,7 +19,7 @@ metadata = {
@fixture @fixture
@patch('orchestrator.activities.slot_manager.Redis.__init__') @patch('orchestrator.activities.slot_manager.RedisRepository')
def slot_manager(_redis_mock): def slot_manager(_redis_mock):
slot_manager = SlotManager( slot_manager = SlotManager(
host='localhost', host='localhost',
@@ -28,31 +28,36 @@ def slot_manager(_redis_mock):
password='password', password='password',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
slot_manager.redis_client = MagicMock() slot_manager.redis_repository = MagicMock()
slot_manager.logger = MagicMock() slot_manager.logger = MagicMock()
slot_manager.notification_handler = MagicMock() slot_manager.notification_handler = MagicMock()
slot_manager.send_notification = MagicMock() slot_manager.send_notification = MagicMock()
slot_manager.send_notification_async = AsyncMock()
slot_manager.emit_metric = AsyncMock()
return slot_manager return slot_manager
@mark.asyncio @mark.asyncio
async def test_load_opc_slots_no_slot_keys(slot_manager): async def test_load_opc_slots_no_slot_keys(slot_manager):
slot_manager.redis_client.keys.return_value = [] slot_manager.redis_repository.keys = AsyncMock(return_value=[])
assert await slot_manager.load_opc_slots(metadata) == {} assert await slot_manager.load_opc_slots(metadata) == {}
@mark.asyncio @mark.asyncio
async def test_load_opc_slots(slot_manager): async def test_load_opc_slots(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_repository.keys = AsyncMock(
b'slot:opc_tags:1', return_value=[
b'slot:opc_tags:2', b'slot:opc_tags:1',
b'slot:opc_tags:3', b'slot:opc_tags:2',
] b'slot:opc_tags:3',
]
)
slot_manager.get = MagicMock(side_effect=['value1', 'value2', None]) slot_manager.redis_repository.get = AsyncMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata) response = await slot_manager.load_opc_slots(metadata)
@@ -65,13 +70,15 @@ async def test_load_opc_slots(slot_manager):
@mark.asyncio @mark.asyncio
async def test_load_opc_slots_no_decode(slot_manager): async def test_load_opc_slots_no_decode(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_repository.keys = AsyncMock(
'slot:opc_tags:1', return_value=[
'slot:opc_tags:2', 'slot:opc_tags:1',
'slot:opc_tags:3', 'slot:opc_tags:2',
] 'slot:opc_tags:3',
]
)
slot_manager.get = MagicMock(side_effect=['value1', 'value2', None]) slot_manager.redis_repository.get = AsyncMock(side_effect=['value1', 'value2', None])
response = await slot_manager.load_opc_slots(metadata) response = await slot_manager.load_opc_slots(metadata)
@@ -84,19 +91,21 @@ async def test_load_opc_slots_no_decode(slot_manager):
@mark.asyncio @mark.asyncio
async def test_load_opc_slots_error(slot_manager): async def test_load_opc_slots_error(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_repository.keys = AsyncMock(
'slot:opc_tags:1', return_value=[
'slot:opc_tags:2', 'slot:opc_tags:1',
'slot:opc_tags:3', 'slot:opc_tags:2',
] 'slot:opc_tags:3',
]
)
slot_manager.get = MagicMock(side_effect=Exception('Test exception')) slot_manager.redis_repository.get = AsyncMock(side_effect=Exception('Test exception'))
try: try:
await slot_manager.load_opc_slots(metadata) await slot_manager.load_opc_slots(metadata)
except Exception as e: except Exception as e:
assert str(e) == 'Test exception' assert str(e) == 'Test exception'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message='Failed to load OPC slots: Test exception', message='Failed to load OPC slots: Test exception',
@@ -111,11 +120,13 @@ async def test_load_opc_slots_error(slot_manager):
@mark.asyncio @mark.asyncio
async def test_load_active_ingestors(slot_manager): async def test_load_active_ingestors(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_repository.keys = AsyncMock(
b'heartbeat:ingestor:1', return_value=[
b'heartbeat:ingestor:2', b'heartbeat:ingestor:1',
'heartbeat:ingestor:3', b'heartbeat:ingestor:2',
] 'heartbeat:ingestor:3',
]
)
response = await slot_manager.load_active_ingestors(metadata) response = await slot_manager.load_active_ingestors(metadata)
@@ -124,19 +135,21 @@ async def test_load_active_ingestors(slot_manager):
@mark.asyncio @mark.asyncio
async def test_load_active_ingestors_error(slot_manager): async def test_load_active_ingestors_error(slot_manager):
slot_manager.redis_client.keys.return_value = [ slot_manager.redis_repository.keys = AsyncMock(
'heartbeat:ingestor:1', return_value=[
'heartbeat:ingestor:2', 'heartbeat:ingestor:1',
'heartbeat:ingestor:3', 'heartbeat:ingestor:2',
] 'heartbeat:ingestor:3',
]
)
slot_manager.redis_client.keys.side_effect = Exception('Test exception') slot_manager.redis_repository.keys = AsyncMock(side_effect=Exception('Test exception'))
try: try:
await slot_manager.load_active_ingestors(metadata) await slot_manager.load_active_ingestors(metadata)
except Exception as e: except Exception as e:
assert str(e) == 'Test exception' assert str(e) == 'Test exception'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message='Failed to load active ingestors: Test exception', message='Failed to load active ingestors: Test exception',
@@ -151,11 +164,11 @@ async def test_load_active_ingestors_error(slot_manager):
@mark.asyncio @mark.asyncio
async def test_update_slots(slot_manager): async def test_update_slots(slot_manager):
slot_manager.set = MagicMock(side_effect=[None, Exception('Test exception')]) slot_manager.redis_repository.set = AsyncMock(side_effect=[None, Exception('Test exception')])
response = await slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}}) response = await slot_manager.update_slots({'to_insert': {'1': 'value1', '2': 'value2'}})
slot_manager.set.assert_has_calls( slot_manager.redis_repository.set.assert_has_calls(
[call('slot:opc_tags:1', 'value1', ttl=None), call('slot:opc_tags:2', 'value2', ttl=None)] [call('slot:opc_tags:1', 'value1', ttl=None), call('slot:opc_tags:2', 'value2', ttl=None)]
) )
@@ -167,11 +180,13 @@ async def test_update_slots(slot_manager):
@mark.asyncio @mark.asyncio
async def test_delete_slots(slot_manager): async def test_delete_slots(slot_manager):
slot_manager.redis_client.delete = MagicMock(side_effect=[None, Exception('Test exception')]) slot_manager.redis_repository.delete = AsyncMock(
side_effect=[None, Exception('Test exception')]
)
response = await slot_manager.delete_slots({'to_delete': ['1', '2']}) response = await slot_manager.delete_slots({'to_delete': ['1', '2']})
slot_manager.redis_client.delete.assert_has_calls( slot_manager.redis_repository.delete.assert_has_calls(
[call('slot:opc_tags:1'), call('slot:opc_tags:2')] [call('slot:opc_tags:1'), call('slot:opc_tags:2')]
) )
@@ -191,7 +206,7 @@ async def test_get_last_data_timestamp_none(slot_manager):
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
} }
slot_manager.get = MagicMock(return_value=None) slot_manager.redis_repository.get = AsyncMock(return_value=None)
result = await slot_manager.get_last_data_timestamp(test_data) result = await slot_manager.get_last_data_timestamp(test_data)
@@ -208,11 +223,13 @@ async def test_get_last_data_timestamp_not_none(slot_manager):
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
} }
slot_manager.get = MagicMock(return_value='2023-01-01 12:00:00') slot_manager.redis_repository.get = AsyncMock(return_value='2023-01-01 12:00:00')
result = await slot_manager.get_last_data_timestamp(test_data) result = await slot_manager.get_last_data_timestamp(test_data)
slot_manager.get.assert_called_once_with('notification_last_timestamp:test_mail_type') slot_manager.redis_repository.get.assert_called_once_with(
'notification_last_timestamp:test_mail_type'
)
assert result == '2023-01-01 12:00:00' assert result == '2023-01-01 12:00:00'
@@ -227,8 +244,8 @@ async def test_get_last_data_timestamp_error(slot_manager):
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
} }
slot_manager.send_notification = MagicMock() slot_manager.send_notification_async = AsyncMock()
slot_manager.get = MagicMock(side_effect=Exception('test')) slot_manager.redis_repository.get = AsyncMock(side_effect=Exception('test'))
try: try:
await slot_manager.get_last_data_timestamp(test_data) await slot_manager.get_last_data_timestamp(test_data)
@@ -236,7 +253,7 @@ async def test_get_last_data_timestamp_error(slot_manager):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_GET_ERROR', notification_id='REDIS_GET_ERROR',
message='Error getting last data timestamp: test', message='Error getting last data timestamp: test',
@@ -288,13 +305,13 @@ async def test_put_last_data_timestamp_not_empty_dataframe(slot_manager):
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
} }
slot_manager.set = MagicMock() slot_manager.redis_repository.set = AsyncMock()
result = await slot_manager.put_last_data_timestamp(test_data) result = await slot_manager.put_last_data_timestamp(test_data)
assert result == '2023-01-01 12:00:01' assert result == '2023-01-01 12:00:01'
slot_manager.set.assert_called_once_with( slot_manager.redis_repository.set.assert_called_once_with(
'notification_last_timestamp:test_mail_type', '2023-01-01 12:00:01', ttl=18000 'notification_last_timestamp:test_mail_type', '2023-01-01 12:00:01', ttl=18000
) )
@@ -316,8 +333,8 @@ async def test_put_last_data_timestamp_error(slot_manager):
'mail_type': 'test_mail_type', 'mail_type': 'test_mail_type',
} }
slot_manager.send_notification = MagicMock() slot_manager.send_notification_async = AsyncMock()
slot_manager.set = MagicMock(side_effect=Exception('test')) slot_manager.redis_repository.set = AsyncMock(side_effect=Exception('test'))
try: try:
await slot_manager.put_last_data_timestamp(test_data) await slot_manager.put_last_data_timestamp(test_data)
@@ -325,7 +342,7 @@ async def test_put_last_data_timestamp_error(slot_manager):
except Exception as e: except Exception as e:
assert str(e) == 'test' assert str(e) == 'test'
slot_manager.send_notification.assert_called_once_with( slot_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='REDIS_SET_ERROR', notification_id='REDIS_SET_ERROR',
message='Error setting last data timestamp: test', message='Error setting last data timestamp: test',
@@ -340,7 +357,7 @@ async def test_put_last_data_timestamp_error(slot_manager):
@mark.asyncio @mark.asyncio
async def test_filter_notification_alerts(slot_manager): async def test_filter_notification_alerts(slot_manager):
slot_manager.get = MagicMock( slot_manager.redis_repository.get = AsyncMock(
side_effect=[ side_effect=[
None, None,
(now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ), (now() - timedelta(seconds=600)).strftime(DATETIME_FORMAT_MS_WITH_TZ),
@@ -402,8 +419,10 @@ async def test_store_notification_cache(slot_manager):
'sent_ttl': 600, 'sent_ttl': 600,
} }
slot_manager.set = MagicMock() slot_manager.redis_repository.set = AsyncMock()
await slot_manager.store_notification_cache(test_data) await slot_manager.store_notification_cache(test_data)
slot_manager.set.assert_called_once_with('test_schedule_1:test_notification_id_1', ANY, ttl=600) slot_manager.redis_repository.set.assert_called_once_with(
'test_schedule_1:test_notification_id_1', ANY, ttl=600
)

View File

@@ -26,10 +26,13 @@ def temporal_manager(connect_mock):
laborious_namespace='laborious', laborious_namespace='laborious',
logger=MagicMock(), logger=MagicMock(),
notification_handler=MagicMock(), notification_handler=MagicMock(),
metrics_controller=AsyncMock(),
) )
temporal_manager.temporal_clients['scouter'] = MagicMock() temporal_manager.temporal_clients['scouter'] = MagicMock()
temporal_manager.temporal_clients['laborious'] = MagicMock() temporal_manager.temporal_clients['laborious'] = MagicMock()
temporal_manager.send_notification_async = AsyncMock()
temporal_manager.emit_metric = AsyncMock()
temporal_manager.send_notification = MagicMock() temporal_manager.send_notification = MagicMock()
return temporal_manager return temporal_manager
@@ -100,7 +103,7 @@ async def test_normalize_schedules_error(temporal_manager):
await temporal_manager.normalize_schedules(metadata) await temporal_manager.normalize_schedules(metadata)
except Exception as e: except Exception as e:
assert str(e) == 'Test exception' assert str(e) == 'Test exception'
temporal_manager.send_notification.assert_called_once_with( temporal_manager.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'], metadata=metadata['metadata'],
notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR', notification_id='SCHEDULER_NORMALIZE_SCHEDULES_ERROR',
message='Failed to normalize schedules: Test exception', message='Failed to normalize schedules: Test exception',

View File

@@ -171,13 +171,6 @@ env:
name: redis name: redis
key: redis-password key: redis-password
- name: COUCHBASE_CONNECTION_STRING
value: "couchbase://sientia.couchbase.svc.cluster.local"
- name: COUCHBASE_USERNAME
value: "sientia"
- name: COUCHBASE_PASSWORD
value: "sientia"
- name: MONGODB_USERNAME - name: MONGODB_USERNAME
value: "root" value: "root"
- name: MONGODB_PASSWORD - name: MONGODB_PASSWORD