From f7b73618d8edfcd3bc20aee8e5079ece7d4b7275 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 16 Jun 2025 16:01:20 -0300 Subject: [PATCH] SIENTIAPDE-1107 feat: implement MongoDB integration with clear_mongo_id function and related tests --- orchestrator/activities/mongo_db.py | 35 +-- .../activities/test_activities.py | 40 ++- .../orchestrator/activities/test_couchbase.py | 12 + .../orchestrator/activities/test_mongo_db.py | 246 ++++++++++++++++++ .../utils/test_connectors_config.py | 27 +- .../workflows/test_orchestrator.py | 4 +- 6 files changed, 332 insertions(+), 32 deletions(-) create mode 100644 tests/orchestrator/activities/test_mongo_db.py diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 866e3b9..e5c663f 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -1,10 +1,10 @@ -import pymongo from temporalio import workflow, activity with workflow.unsafe.imports_passed_through(): from typing import Any import traceback from logging import Logger + from pymongo import MongoClient from sientia_do.notifications.handlers import NotificationHandler from sientia_do.notifications.models import NotificationLevel from sientia_do.temporal.activities.base import BaseActivity @@ -21,25 +21,18 @@ def clear_mongo_id(docs: list) -> list: list: The documents without the `_id` field. """ for doc in docs: - if not isinstance(doc, dict) and not isinstance(doc, list): - continue - - if "_id" in doc: - del doc["_id"] - - if isinstance(doc, dict): - for _key, value in doc.items(): - if isinstance(value, dict): - clear_mongo_id([value]) - elif isinstance(value, list): - clear_mongo_id(value) - if isinstance(doc, list): - for item in doc: - if isinstance(item, dict): - clear_mongo_id([item]) - elif isinstance(item, list): - clear_mongo_id(item) + 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 @@ -51,14 +44,14 @@ class MongoDB(BaseActivity): self.connection_string = connection_string self.database_name = database_name - self.client = pymongo.MongoClient( + self.client = MongoClient( self.connection_string, serverSelectionTimeoutMS=5000) self.client.server_info() # Trigger an exception if connection fails self.database = self.client[self.database_name] # Initialize MongoDB client here (omitted for brevity) - self.logger.info("MongoDB connection initialized") + logger.info("MongoDB connection initialized") BaseActivity.__init__(self, logger=logger, diff --git a/tests/orchestrator/activities/test_activities.py b/tests/orchestrator/activities/test_activities.py index 997a55a..7106085 100644 --- a/tests/orchestrator/activities/test_activities.py +++ b/tests/orchestrator/activities/test_activities.py @@ -1,18 +1,21 @@ from unittest.mock import patch, MagicMock, ANY from pytest import mark +from orchestrator.activities import mongo_db from orchestrator.activities.activities import Activities -from orchestrator.activities.couchbase import Couchbase +from orchestrator.activities.mongo_db import MongoDB from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.activities.slot_manager import SlotManager from orchestrator.activities.formatters import Formatters @patch('orchestrator.activities.couchbase.Couchbase.__init__') +@patch('orchestrator.activities.mongo_db.MongoDB.__init__') @patch('orchestrator.activities.temporal_manager.TemporalManager.__init__') @patch('orchestrator.activities.slot_manager.SlotManager.__init__') @patch('orchestrator.activities.formatters.Formatters.__init__') def test___init__(mock_formatters_init, mock_slot_manager_init, mock_temporal_manager_init, + mock_mongodb_init, mock_couchbase_init): couchbase_config = { @@ -21,6 +24,11 @@ def test___init__(mock_formatters_init, mock_slot_manager_init, 'password': 'password' } + mongo_db_config = { + 'connection_string': 'mongodb://localhost:27017', + 'database_name': 'test_db' + } + redis_config = { 'host': 'localhost', 'port': 6379, @@ -34,14 +42,16 @@ def test___init__(mock_formatters_init, mock_slot_manager_init, activities = Activities( temporal_client=temporal_client, - couchbase_config=couchbase_config, + # couchbase_config=couchbase_config, redis_config=redis_config, + mongodb_config=mongo_db_config, logger=logger, notification_handler=notification_handler ) assert isinstance(activities, Activities) - assert isinstance(activities, Couchbase) + # assert isinstance(activities, Couchbase) + assert isinstance(activities, MongoDB) assert isinstance(activities, TemporalManager) assert isinstance(activities, SlotManager) assert isinstance(activities, Formatters) @@ -56,11 +66,19 @@ def test___init__(mock_formatters_init, mock_slot_manager_init, notification_handler=notification_handler ) - mock_couchbase_init.assert_called_once_with( + # mock_couchbase_init.assert_called_once_with( + # ANY, + # connection_string=couchbase_config['connection_string'], + # username=couchbase_config['username'], + # password=couchbase_config['password'], + # logger=logger, + # notification_handler=notification_handler + # ) + + mock_mongodb_init.assert_called_once_with( ANY, - connection_string=couchbase_config['connection_string'], - username=couchbase_config['username'], - password=couchbase_config['password'], + connection_string='mongodb://localhost:27017', + database_name='test_db', logger=logger, notification_handler=notification_handler ) @@ -88,6 +106,11 @@ async def test_prepare_activity(_mock_cluster): 'password': 'password' } + mongo_db_config = { + 'connection_string': 'mongodb://localhost:27017', + 'database_name': 'test_db' + } + redis_config = { 'host': 'localhost', 'port': 6379, @@ -101,8 +124,9 @@ async def test_prepare_activity(_mock_cluster): activities = Activities( temporal_client=temporal_client, - couchbase_config=couchbase_config, + # couchbase_config=couchbase_config, redis_config=redis_config, + mongodb_config=mongo_db_config, logger=logger, notification_handler=notification_handler ) diff --git a/tests/orchestrator/activities/test_couchbase.py b/tests/orchestrator/activities/test_couchbase.py index 84f4da7..91b418a 100644 --- a/tests/orchestrator/activities/test_couchbase.py +++ b/tests/orchestrator/activities/test_couchbase.py @@ -16,6 +16,18 @@ def couchbase(_cluster_mock): ) +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( + "Failed to close Couchbase connection: %s", couchbase.cluster.close.side_effect) + + @mark.asyncio async def test_load_query_from_couchbase_success(couchbase): couchbase.cluster.query.return_value.rows.return_value = [ diff --git a/tests/orchestrator/activities/test_mongo_db.py b/tests/orchestrator/activities/test_mongo_db.py new file mode 100644 index 0000000..cb89784 --- /dev/null +++ b/tests/orchestrator/activities/test_mongo_db.py @@ -0,0 +1,246 @@ +from unittest.mock import MagicMock, patch, ANY +from pytest import fixture, mark +from orchestrator.activities.mongo_db import clear_mongo_id +from orchestrator.activities.mongo_db import MongoDB +from sientia_do.notifications.models import NotificationLevel + + +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 +@patch("orchestrator.activities.mongo_db.MongoClient") +def mongo_db(mongo_mock): + return ( + MongoDB( + connection_string="mongodb://localhost:27017", + database_name="test_db", + logger=MagicMock(), + notification_handler=MagicMock() + ) + ) + + +@patch("orchestrator.activities.mongo_db.MongoClient") +def test___init__(mongo_mock): + mongo_db = MongoDB( + connection_string="mongodb://localhost:27017", + database_name="test_db", + logger=MagicMock(), + notification_handler=MagicMock() + ) + 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") + + +def test_shutdown_success(mongo_db): + mongo_db.shutdown() + mongo_db.client.close.assert_called_once() + mongo_db.logger.info.assert_any_call("Closing MongoDB connection...") + mongo_db.logger.info.assert_any_call( + "MongoDB connection closed successfully") + + +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" + ) + + +@mark.asyncio +async def test_find_documents_in_mongodb_success(mongo_db): + input_data = {"collection": "test_collection", "filters": { + "name": {"$exists": True} + }} + mock_collection = MagicMock() + mock_collection.find.return_value = [ + {"_id": "12345", "name": "test1"}, + {"_id": "67890", "name": "test2"} + ] + mongo_db.database.__getitem__.return_value = mock_collection + + result = await mongo_db.find_documents_in_mongodb( + { + "query": input_data + }) + + assert len(result) == 2 + assert result[0] == {"name": "test1"} + assert result[1] == {"name": "test2"} + mock_collection.find.assert_called_once_with( + {"name": {"$exists": True}}, {"_id": 0} + ) + + +@mark.asyncio +async def test_find_documents_in_mongodb_failure(mongo_db): + input_data = {"collection": "test_collection", "filters": { + "name": {"$exists": True} + }} + mongo_db.database.__getitem__.return_value = MagicMock( + find=MagicMock(side_effect=Exception("Error")) + ) + + try: + await mongo_db.find_documents_in_mongodb( + { + "query": input_data + }) + except Exception as e: + assert str(e) == "Error" + mongo_db.notification_handler.build_and_send_notification.assert_called_once_with( + notification_id="MONGODB_QUERY_ERROR", + message="Failed to execute MongoDB query: Error", + level=NotificationLevel.ERROR, + block="load_query_from_mongodb", + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + +@mark.asyncio +async def test_find_documents_in_mongodb_missing_collection(mongo_db): + input_data = {"query": {"filters": {}}} + + try: + await mongo_db.find_documents_in_mongodb(input_data) + + except ValueError as e: + assert str(e) == "Collection name must be provided in the query." + + else: + assert False, "Expected a ValueError to be raised" + + +@mark.asyncio +async def test_aggregate_documents_in_mongodb_success(mongo_db): + input_data = {"collection": "test_collection", "aggregation": [ + {"$match": {"name": {"$exists": True}}}, + {"$project": {"name": 1}} + ]} + mock_collection = MagicMock() + mock_collection.aggregate.return_value = [ + {"_id": "asdad", "name": "test1"}, + {"_id": "adzx", "name": "test2"} + ] + mongo_db.database.__getitem__.return_value = mock_collection + + result = await mongo_db.aggregate_documents_in_mongodb( + { + "query": input_data + }) + + assert len(result) == 2 + assert result[0] == {"name": "test1"} + assert result[1] == {"name": "test2"} + expected_pipeline = input_data["aggregation"] + expected_pipeline.append({"$project": {"_id": 0}}) + + mock_collection.aggregate.assert_called_once_with( + expected_pipeline + ) + + +@mark.asyncio +async def test_aggregate_documents_in_mongodb_failure(mongo_db): + input_data = {"collection": "test_collection", "aggregation": [ + {"$match": {"name": {"$exists": True}}}, + {"$project": {"name": 1}} + ]} + mongo_db.database.__getitem__.return_value = MagicMock( + aggregate=MagicMock(side_effect=Exception("Error")) + ) + + try: + await mongo_db.aggregate_documents_in_mongodb( + { + "query": input_data + }) + except Exception as e: + assert str(e) == "Error" + mongo_db.notification_handler.build_and_send_notification.assert_called_once_with( + notification_id="MONGODB_AGGREGATION_ERROR", + message="Failed to execute MongoDB aggregation: Error", + level=NotificationLevel.ERROR, + block="aggregate_documents_in_mongodb", + attachment_content=ANY + ) + + else: + assert False, "Expected an exception to be raised" + + +@mark.asyncio +async def test_aggregate_documents_in_mongodb_missing_collection(mongo_db): + input_data = {"query": {"aggregation": []}} + + try: + await mongo_db.aggregate_documents_in_mongodb(input_data) + + except ValueError as e: + assert str(e) == "Collection name must be provided in the query." + + else: + assert False, "Expected a ValueError to be raised" + + +@mark.asyncio +async def test_aggregate_documents_in_mongodb_missing_aggregation(mongo_db): + input_data = {"query": {"collection": "test_collection"}} + + try: + await mongo_db.aggregate_documents_in_mongodb(input_data) + + except ValueError as e: + assert str(e) == "Aggregation must be provided." + + else: + assert False, "Expected a ValueError to be raised" diff --git a/tests/orchestrator/utils/test_connectors_config.py b/tests/orchestrator/utils/test_connectors_config.py index 3cf4dac..06b1606 100644 --- a/tests/orchestrator/utils/test_connectors_config.py +++ b/tests/orchestrator/utils/test_connectors_config.py @@ -1,6 +1,7 @@ from os import environ from orchestrator.utils.connectors_config import (build_redis_config, - build_couchbase_config) + build_couchbase_config, + build_mongodb_config) def test_build_redis_config_with_env_vars(): @@ -49,3 +50,27 @@ def test_build_couchbase_config_with_defaults(): 'username': 'sientia', 'password': 'sientia' } + + +def test_build_mongo_db_config_with_env_vars(): + environ['MONGODB_USERNAME'] = 'sientia1' + environ['MONGODB_PASSWORD'] = 'sientia1' + environ['MONGODB_URL'] = 'localhost:27018' + environ['MONGODB_DATABASE_NAME'] = 'test_db' + + assert build_mongodb_config() == { + 'connection_string': 'mongodb://sientia1:sientia1@localhost:27018', + 'database_name': 'test_db' + } + + +def test_build_mongo_db_config_with_defaults(): + environ.pop('MONGODB_USERNAME', None) + environ.pop('MONGODB_PASSWORD', None) + environ.pop('MONGODB_DATABASE_NAME', None) + environ.pop('MONGODB_URL', None) + + assert build_mongodb_config() == { + 'connection_string': 'mongodb://sientia:sientia@localhost:27017', + 'database_name': 'sientia' + } diff --git a/tests/orchestrator/workflows/test_orchestrator.py b/tests/orchestrator/workflows/test_orchestrator.py index 0b1d4f8..3a131a4 100644 --- a/tests/orchestrator/workflows/test_orchestrator.py +++ b/tests/orchestrator/workflows/test_orchestrator.py @@ -36,7 +36,7 @@ async def test_run(workflow_mock, orchestrator): workflow_mock.execute_local_activity_method.assert_has_calls([ call( - Activities.load_query_from_couchbase, + Activities.aggregate_documents_in_mongodb, { "query": input_data["pipelines_query"] }, @@ -47,7 +47,7 @@ async def test_run(workflow_mock, orchestrator): workflow_mock.execute_local_activity_method.assert_has_calls([ call( - Activities.load_query_from_couchbase, + Activities.find_documents_in_mongodb, { "query": input_data["opc_servers_query"] },