gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1107

Browse Source 
feat: implement MongoDB integration with clear_mongo_id function and related tests
This commit is contained in:
vitor-aignosi
2025-06-16 16:01:20 -03:00
parent 7984cb29f4
commit f7b73618d8
6 changed files with 332 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
orchestrator/activities/mongo_db.py
View File

@@ -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 isinstance(doc, list):
clear_mongo_id(doc)
elif isinstance(doc, dict):
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):
for key, value in doc.items():
if 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)
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,

40
tests/orchestrator/activities/test_activities.py
View File

@@ -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
)

12
tests/orchestrator/activities/test_couchbase.py
View File

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

246
tests/orchestrator/activities/test_mongo_db.py Normal file
View File

@@ -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"

27
tests/orchestrator/utils/test_connectors_config.py
View File

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

4
tests/orchestrator/workflows/test_orchestrator.py
View File

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