sientia-dataops-scouter_temporal/e2e/fixtures/fake_mongodb_repository.py

"""
Fake MongoDB Repository adapter for testing.

This adapter implements the MongoDBRepository interface using mongomock
to provide an in-memory MongoDB server for testing.
"""

import time
from typing import Any

from mongomock import MongoClient
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring


def clear_mongo_id(docs: list) -> list:
    """
    Remove MongoDB internal `_id` fields from nested structures.

    Same implementation as in the real MongoDBRepository.

    Args:
        docs: The list of documents or nested structures to clean.

    Returns:
        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 FakeMongoDBRepository(SientiaMonitoring):
    """
    Fake MongoDB Repository that uses mongomock for testing.

    Implements the same interface as MongoDBRepository but uses
    mongomock for in-memory MongoDB operations.
    """

    def __init__(
        self,
        connection_string: str,
        database_name: str,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        """
        Initialize fake MongoDB repository with mongomock.

        Args:
            connection_string: MongoDB connection string (ignored in fake mode)
            database_name: Target database name
            logger: Logger instance
            notification_handler: Notification handler
            metrics_controller: Metrics controller
        """
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=metrics_controller,
        )

        self.database_name = database_name
        # Use mongomock instead of real MongoDB
        self.mongo_client = MongoClient()
        self.database = self.mongo_client[self.database_name]
        logger.info('Fake MongoDB connection initialized')

    def close(self):
        """
        Closes fake MongoDB connection and shuts down monitoring.
        """
        try:
            if self.mongo_client:
                self.logger.info('Closing fake MongoDB connection...')
                self.mongo_client.close()
                self.logger.info('Fake MongoDB connection closed successfully')
        except Exception as e:
            self.logger.error(f'Failed to close fake MongoDB connection: {e}')
        SientiaMonitoring.shutdown(self)

    def __del__(self):
        """
        Destructor that closes the connection when destroying the instance.
        """
        self.close()

    async def find(
        self,
        collection_name: str,
        filters: dict[str, Any],
        metadata: dict[str, Any],
    ) -> list[dict[str, Any]]:
        """
        Finds documents in a MongoDB collection based on provided filters.

        Args:
            collection_name: Name of the collection to search
            filters: Query filters to apply
            metadata: Dictionary with additional metadata

        Return:
            List of documents matching the filters (with `_id` removed)
        """
        try:
            start_time = time.time()
            collection = self.database[collection_name]
            documents = list(collection.find(filters, {'_id': 0}))
        except Exception as e:
            self.logger.error(f'Failed to find documents in fake MongoDB: {e}')
            raise e

        return clear_mongo_id(documents)

    async def aggregate(
        self,
        collection_name: str,
        pipeline: list[dict[str, Any]],
        metadata: dict[str, Any],
    ) -> list[dict[str, Any]]:
        """
        Executes an aggregation on a MongoDB collection.

        Args:
            collection_name: Name of the collection to aggregate
            pipeline: MongoDB aggregation pipeline
            metadata: Dictionary with additional metadata

        Return:
            List of documents resulting from aggregation (with `_id` removed)
        """
        try:
            start_time = time.time()
            collection = self.database[collection_name]
            documents = list(collection.aggregate(pipeline))
        except Exception as e:
            self.logger.error(f'Failed to aggregate documents in fake MongoDB: {e}')
            raise e

        return clear_mongo_id(documents)

    async def update_many(
        self,
        collection_name: str,
        filters: dict[str, Any],
        update: dict[str, Any],
        metadata: dict[str, Any],
    ) -> None:
        """
        Updates multiple documents in a MongoDB collection.

        Args:
            collection_name: Collection name
            filters: Filters to identify documents to update
            update: Update operations to apply
            metadata: Dictionary with additional metadata
        """
        try:
            collection = self.database[collection_name]
            collection.update_many(filters, update)
        except Exception as e:
            self.logger.error(f'Failed to update documents in fake MongoDB: {e}')
            raise e

    async def insert_many(
        self,
        collection_name: str,
        documents: list[dict[str, Any]],
        metadata: dict[str, Any],
    ) -> None:
        """
        Inserts multiple documents into a MongoDB collection.

        Args:
            collection_name: Collection name
            documents: List of documents to insert
            metadata: Dictionary with additional metadata
        """
        try:
            collection = self.database[collection_name]
            collection.insert_many(documents)
        except Exception as e:
            self.logger.error(f'Failed to insert documents in fake MongoDB: {e}')
            raise e

    async def insert(
        self,
        collection_name: str,
        document: dict[str, Any],
        metadata: dict[str, Any],
    ) -> None:
        """
        Inserts a document into a MongoDB collection.

        Args:
            collection_name: Collection name
            document: Document to insert
            metadata: Dictionary with additional metadata
        """
        try:
            collection = self.database[collection_name]
            collection.insert_one(document)
        except Exception as e:
            self.logger.error(f'Failed to insert document in fake MongoDB: {e}')
            raise e

    async def delete_many(
        self,
        collection_name: str,
        filters: dict[str, Any],
        metadata: dict[str, Any],
    ) -> None:
        """
        Removes multiple documents from a MongoDB collection.

        Args:
            collection_name: Collection name
            filters: Filters to identify documents to remove
            metadata: Dictionary with additional metadata
        """
        try:
            collection = self.database[collection_name]
            collection.delete_many(filters)
        except Exception as e:
            self.logger.error(f'Failed to delete documents in fake MongoDB: {e}')
            raise e