From 9ecc8a8afee4c8dec2c993100e2de7181f6b8196 Mon Sep 17 00:00:00 2001
From: Akash Verma <138790903+akashverma0786@users.noreply.github.com>
Date: Fri, 10 Jan 2025 10:10:11 +0530
Subject: [PATCH] Added integration testcontainer test for mongodb (#19282)

---
 ingestion/tests/integration/conftest.py       |  9 ++-
 .../tests/integration/mongodb/conftest.py     | 59 ++++++++++++++
 .../integration/mongodb/test_metadata.py      | 78 +++++++++++++++++++
 3 files changed, 144 insertions(+), 2 deletions(-)
 create mode 100644 ingestion/tests/integration/mongodb/conftest.py
 create mode 100644 ingestion/tests/integration/mongodb/test_metadata.py

diff --git a/ingestion/tests/integration/conftest.py b/ingestion/tests/integration/conftest.py
index 081160c303c..a838715059a 100644
--- a/ingestion/tests/integration/conftest.py
+++ b/ingestion/tests/integration/conftest.py
@@ -127,8 +127,13 @@ def unmask_password(create_service_request):
     """
 
     def patch_password(service: DatabaseService):
-        service.connection.config.authType.password = (
-            create_service_request.connection.config.authType.password
+        if hasattr(service.connection.config, "authType"):
+            service.connection.config.authType.password = (
+                create_service_request.connection.config.authType.password
+            )
+            return service
+        service.connection.config.password = (
+            create_service_request.connection.config.password
         )
         return service
 
diff --git a/ingestion/tests/integration/mongodb/conftest.py b/ingestion/tests/integration/mongodb/conftest.py
new file mode 100644
index 00000000000..6b0532d4f61
--- /dev/null
+++ b/ingestion/tests/integration/mongodb/conftest.py
@@ -0,0 +1,59 @@
+import os
+from datetime import datetime
+
+import pytest
+
+from _openmetadata_testutils.helpers.docker import try_bind
+from metadata.generated.schema.api.services.createDatabaseService import (
+    CreateDatabaseServiceRequest,
+)
+from metadata.generated.schema.entity.services.connections.database.mongoDBConnection import (
+    MongoDBConnection,
+)
+from metadata.generated.schema.entity.services.databaseService import (
+    DatabaseConnection,
+    DatabaseServiceType,
+)
+
+
+@pytest.fixture(scope="module")
+def mongodbContainer(tmp_path_factory):
+    """
+    Start a Mongodb container
+    """
+    from testcontainers.mongodb import MongoDbContainer
+
+    container = MongoDbContainer()
+
+    with (
+        try_bind(container, 27017, None) if not os.getenv("CI") else container
+    ) as container:
+        db = container.get_connection_client().test
+        db.user_profiles.insert_one(
+            {
+                "user_id": "123e4567-e89b-12d3-a456-426614174000",
+                "first_name": "John",
+                "last_name": "Doe",
+                "email": "john.doe@example.com",
+                "signup_date": datetime.now(),
+                "is_active": True,
+            }
+        )
+
+        yield container
+
+
+@pytest.fixture(scope="module")
+def create_service_request(mongodbContainer, tmp_path_factory):
+    return CreateDatabaseServiceRequest(
+        name="docker_test_" + tmp_path_factory.mktemp("mongodb").name,
+        serviceType=DatabaseServiceType.MongoDB,
+        connection=DatabaseConnection(
+            config=MongoDBConnection(
+                username="test",
+                password="test",
+                hostPort=f"localhost:{mongodbContainer.get_exposed_port(27017)}",
+                databaseName="test",
+            )
+        ),
+    )
diff --git a/ingestion/tests/integration/mongodb/test_metadata.py b/ingestion/tests/integration/mongodb/test_metadata.py
new file mode 100644
index 00000000000..620ec3476ff
--- /dev/null
+++ b/ingestion/tests/integration/mongodb/test_metadata.py
@@ -0,0 +1,78 @@
+import sys
+
+import pytest
+
+from metadata.generated.schema.entity.data.table import Table
+from metadata.workflow.metadata import MetadataWorkflow
+
+if not sys.version_info >= (3, 9):
+    pytest.skip("requires python 3.9+", allow_module_level=True)
+
+
+@pytest.fixture(scope="module")
+def prepare_mongodb(mongodbContainer):
+    db = mongodbContainer.get_connection_client().test
+    db.create_collection(
+        "test_table",
+        validator={
+            "$jsonSchema": {
+                "bsonType": "object",
+                "required": ["id", "name"],
+                "properties": {
+                    "id": {
+                        "bsonType": "int",
+                    },
+                    "name": {
+                        "bsonType": "string",
+                        "maxLength": 100,
+                    },
+                    "age": {
+                        "bsonType": "int",
+                    },
+                },
+            }
+        },
+    )
+    data = [
+        {"id": 1, "name": "John Doe", "age": 25},
+        {"id": 2, "name": "Jane Smith", "age": 30},
+    ]
+
+    db.test_table.insert_many(data)
+
+
+@pytest.mark.parametrize(
+    "table_fqn,expected_columns",
+    [
+        [
+            "{service}.test.test.test_table",
+            {
+                "id": {"type": "int", "nullable": False},
+                "name": {"type": "string", "nullable": False},
+                "age": {"type": "int", "nullable": True},
+            },
+        ]
+    ],
+    ids=lambda x: x.split(".")[-1] if isinstance(x, str) else "",
+)
+def test_ingest_metadata(
+    patch_passwords_for_db_services,
+    run_workflow,
+    ingestion_config,
+    metadata,
+    table_fqn,
+    expected_columns,
+    db_service,
+    prepare_mongodb,
+):
+    run_workflow(MetadataWorkflow, ingestion_config)
+    table = metadata.get_by_name(
+        entity=Table, fqn=table_fqn.format(service=db_service.fullyQualifiedName.root)
+    )
+    assert table
+    assert table.fullyQualifiedName.root.split(".")[-1] == "test_table"
+    assert len(table.columns) == 4
+    for name, properties in expected_columns.items():
+        column = next((col for col in table.columns if col.name.root == name), None)
+        assert column is not None
+        assert column.dataType.name.lower() == properties["type"]