SIENTIAPDE-1107

feat: implement MongoDB integration with clear_mongo_id function and related tests

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:
+            clear_mongo_id(doc)
-                if isinstance(item, dict):
+
-                    clear_mongo_id([item])
+        elif isinstance(doc, dict):
-                elif isinstance(item, list):
+            if "_id" in doc:
-                    clear_mongo_id(item)
+                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,

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'],
+        connection_string='mongodb://localhost:27017',
-        username=couchbase_config['username'],
+        database_name='test_db',
-        password=couchbase_config['password'],
         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
     )

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 = [

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"

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'
+    }

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"]
             },