This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9c40220598d Remove mongomock dependency from mongo provider and
introduce MongoDbContainer testcontainer for testing. (#53773)
9c40220598d is described below
commit 9c40220598d706f26db62a967ad72fbc017a8946
Author: GPK <[email protected]>
AuthorDate: Sat Jul 26 13:25:45 2025 +0100
Remove mongomock dependency from mongo provider and introduce
MongoDbContainer testcontainer for testing. (#53773)
* Remove mongomock dependency from mongo provider
* Fix mongo integration tests
---
providers/mongo/pyproject.toml | 7 +-
providers/mongo/tests/conftest.py | 9 +
.../mongo/tests/unit/mongo/hooks/test_mongo.py | 249 ++++++++++-----------
scripts/ci/docker-compose/integration-mongo.yml | 4 +-
4 files changed, 132 insertions(+), 137 deletions(-)
diff --git a/providers/mongo/pyproject.toml b/providers/mongo/pyproject.toml
index 40c0f0595b3..39c36474a6c 100644
--- a/providers/mongo/pyproject.toml
+++ b/providers/mongo/pyproject.toml
@@ -59,10 +59,7 @@ requires-python = ">=3.10"
dependencies = [
"apache-airflow>=2.10.0",
"dnspython>=1.13.0",
- # mongomock does not work with pymongo 4.11.
- # See https://github.com/apache/airflow/issues/46215 and this one can be
removed after
- # https://github.com/mongomock/mongomock/issues/912 is closed by mongomock.
- "pymongo>=4.3.0,<4.11.0",
+ "pymongo>=4.13.2",
]
[dependency-groups]
@@ -71,7 +68,7 @@ dev = [
"apache-airflow-task-sdk",
"apache-airflow-devel-common",
# Additional devel dependencies (do not remove this line and add extra
development dependencies)
- "mongomock>=4.3.0",
+ "testcontainers>=4.12.0"
]
# To build docs:
diff --git a/providers/mongo/tests/conftest.py
b/providers/mongo/tests/conftest.py
index f56ccce0a3f..c6a2457e991 100644
--- a/providers/mongo/tests/conftest.py
+++ b/providers/mongo/tests/conftest.py
@@ -16,4 +16,13 @@
# under the License.
from __future__ import annotations
+import pytest
+from testcontainers.mongodb import MongoDbContainer
+
pytest_plugins = "tests_common.pytest_plugin"
+
+
[email protected](scope="session")
+def mongodb_container():
+ with MongoDbContainer("mongo:latest") as mongo:
+ yield mongo.get_connection_url()
diff --git a/providers/mongo/tests/unit/mongo/hooks/test_mongo.py
b/providers/mongo/tests/unit/mongo/hooks/test_mongo.py
index 2122895dea1..02e7b451e42 100644
--- a/providers/mongo/tests/unit/mongo/hooks/test_mongo.py
+++ b/providers/mongo/tests/unit/mongo/hooks/test_mongo.py
@@ -17,11 +17,9 @@
# under the License.
from __future__ import annotations
-import importlib
-from typing import TYPE_CHECKING
-
import pymongo
import pytest
+from pymongo import MongoClient
from pymongo.errors import CollectionInvalid
from airflow.exceptions import AirflowConfigException
@@ -30,20 +28,30 @@ from airflow.providers.mongo.hooks.mongo import MongoHook
from tests_common.test_utils.compat import connection_as_json
-if TYPE_CHECKING:
- from types import ModuleType
-mongomock: ModuleType | None
[email protected]
+def container_mongo_client(mongodb_container):
+ client = MongoClient(mongodb_container)
+
+ yield client
+
+ client.close()
+
+
[email protected]
+def mongo_collection(container_mongo_client):
+ collection = container_mongo_client.test.test_collection
-try:
- mongomock = importlib.import_module("mongomock")
-except ImportError:
- mongomock = None
+ yield collection
+
+ collection.drop()
@pytest.fixture(scope="module", autouse=True)
-def mongo_connections():
+def mongo_connections(mongodb_container):
"""Create MongoDB connections which use for testing purpose."""
+ (mongodb_container_host, mongodb_container_port) =
mongodb_container.split("@")[1].split(":")
+
connections = [
Connection(conn_id="mongo_default", conn_type="mongo", host="mongo",
port=27017),
Connection(
@@ -69,6 +77,15 @@ def mongo_connections():
),
# Mongo establishes connection during initialization, so we need to
have this connection
Connection(conn_id="fake_connection", conn_type="mongo", host="mongo",
port=27017),
+ Connection(
+ conn_id="mongo_default_with_test_container",
+ conn_type="mongo",
+ login="test",
+ password="test",
+ host=mongodb_container_host,
+ port=int(mongodb_container_port),
+ extra='{"allow_insecure": true, "ssl": false}',
+ ),
]
with pytest.MonkeyPatch.context() as mp:
@@ -77,23 +94,9 @@ def mongo_connections():
yield
-class MongoHookTest(MongoHook):
- """
- Extending hook so that a mockmongo collection object can be passed in
- to get_collection()
- """
-
- def __init__(self, mongo_conn_id="mongo_default", *args, **kwargs):
- super().__init__(mongo_conn_id=mongo_conn_id, *args, **kwargs)
-
- def get_collection(self, mock_collection, mongo_db=None):
- return mock_collection
-
-
[email protected](mongomock is None, reason="mongomock package not present")
class TestMongoHook:
def setup_method(self):
- self.hook = MongoHookTest(mongo_conn_id="mongo_default")
+ self.hook =
MongoHook(mongo_conn_id="mongo_default_with_test_container")
self.conn = self.hook.get_conn()
def teardown_method(self):
@@ -106,8 +109,10 @@ class TestMongoHook:
assert MongoHook("fake_connection").mongo_conn_id == "fake_connection"
def test_get_conn(self):
- assert self.hook.connection.port == 27017
- assert isinstance(self.conn, pymongo.MongoClient)
+ hook = MongoHook(mongo_conn_id="mongo_default")
+ conn = hook.get_conn()
+ assert hook.connection.port == 27017
+ assert isinstance(conn, pymongo.MongoClient)
def test_srv(self):
hook = MongoHook(mongo_conn_id="mongo_default_with_srv")
@@ -131,144 +136,130 @@ class TestMongoHook:
with pytest.raises(AirflowConfigException, match="srv URI should not
specify a port"):
MongoHook(mongo_conn_id="mongo_invalid_srv_with_port")
- def test_insert_one(self):
- collection = mongomock.MongoClient().db.collection
+ def test_insert_one(self, mongo_collection):
obj = {"test_insert_one": "test_value"}
- self.hook.insert_one(collection, obj)
+ self.hook.insert_one("test_collection", obj, mongo_db="test")
- result_obj = collection.find_one(filter=obj)
+ result_obj = mongo_collection.find_one(filter=obj)
assert obj == result_obj
- def test_insert_many(self):
- collection = mongomock.MongoClient().db.collection
+ def test_insert_many(self, mongo_collection):
objs = [{"test_insert_many_1": "test_value"}, {"test_insert_many_2":
"test_value"}]
- self.hook.insert_many(collection, objs)
+ self.hook.insert_many("test_collection", objs, mongo_db="test")
- result_objs = list(collection.find())
+ result_objs = list(mongo_collection.find())
assert len(result_objs) == 2
- def test_update_one(self):
- collection = mongomock.MongoClient().db.collection
+ def test_update_one(self, mongo_collection):
obj = {"_id": "1", "field": 0}
- collection.insert_one(obj)
+ mongo_collection.insert_one(obj)
filter_doc = obj
update_doc = {"$inc": {"field": 123}}
- self.hook.update_one(collection, filter_doc, update_doc)
+ self.hook.update_one("test_collection", filter_doc, update_doc,
mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == 123
- def test_update_one_with_upsert(self):
- collection = mongomock.MongoClient().db.collection
-
+ def test_update_one_with_upsert(self, mongo_collection):
filter_doc = {"_id": "1", "field": 0}
update_doc = {"$inc": {"field": 123}}
- self.hook.update_one(collection, filter_doc, update_doc, upsert=True)
+ self.hook.update_one("test_collection", filter_doc, update_doc,
upsert=True, mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == 123
- def test_update_many(self):
- collection = mongomock.MongoClient().db.collection
+ def test_update_many(self, mongo_collection):
obj1 = {"_id": "1", "field": 0}
obj2 = {"_id": "2", "field": 0}
- collection.insert_many([obj1, obj2])
+ mongo_collection.insert_many([obj1, obj2])
filter_doc = {"field": 0}
update_doc = {"$inc": {"field": 123}}
- self.hook.update_many(collection, filter_doc, update_doc)
+ self.hook.update_many("test_collection", filter_doc, update_doc,
mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == 123
- result_obj = collection.find_one(filter="2")
+ result_obj = mongo_collection.find_one(filter="2")
assert result_obj["field"] == 123
- def test_update_many_with_upsert(self):
- collection = mongomock.MongoClient().db.collection
-
+ def test_update_many_with_upsert(self, mongo_collection):
filter_doc = {"_id": "1", "field": 0}
update_doc = {"$inc": {"field": 123}}
- self.hook.update_many(collection, filter_doc, update_doc, upsert=True)
+ self.hook.update_many("test_collection", filter_doc, update_doc,
upsert=True, mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == 123
- def test_replace_one(self):
- collection = mongomock.MongoClient().db.collection
+ def test_replace_one(self, mongo_collection):
obj1 = {"_id": "1", "field": "test_value_1"}
obj2 = {"_id": "2", "field": "test_value_2"}
- collection.insert_many([obj1, obj2])
+ mongo_collection.insert_many([obj1, obj2])
obj1["field"] = "test_value_1_updated"
- self.hook.replace_one(collection, obj1)
+ self.hook.replace_one("test_collection", obj1, mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == "test_value_1_updated"
# Other document should stay intact
- result_obj = collection.find_one(filter="2")
+ result_obj = mongo_collection.find_one(filter="2")
assert result_obj["field"] == "test_value_2"
- def test_replace_one_with_filter(self):
- collection = mongomock.MongoClient().db.collection
+ def test_replace_one_with_filter(self, mongo_collection):
obj1 = {"_id": "1", "field": "test_value_1"}
obj2 = {"_id": "2", "field": "test_value_2"}
- collection.insert_many([obj1, obj2])
+ mongo_collection.insert_many([obj1, obj2])
obj1["field"] = "test_value_1_updated"
- self.hook.replace_one(collection, obj1, {"field": "test_value_1"})
+ self.hook.replace_one("test_collection", obj1, {"field":
"test_value_1"}, mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == "test_value_1_updated"
# Other document should stay intact
- result_obj = collection.find_one(filter="2")
+ result_obj = mongo_collection.find_one(filter="2")
assert result_obj["field"] == "test_value_2"
- def test_replace_one_with_upsert(self):
- collection = mongomock.MongoClient().db.collection
-
+ def test_replace_one_with_upsert(self, mongo_collection):
obj = {"_id": "1", "field": "test_value_1"}
- self.hook.replace_one(collection, obj, upsert=True)
+ self.hook.replace_one("test_collection", obj, upsert=True,
mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == "test_value_1"
- def test_replace_many(self):
- collection = mongomock.MongoClient().db.collection
+ def test_replace_many(self, mongo_collection):
obj1 = {"_id": "1", "field": "test_value_1"}
obj2 = {"_id": "2", "field": "test_value_2"}
- collection.insert_many([obj1, obj2])
+ mongo_collection.insert_many([obj1, obj2])
obj1["field"] = "test_value_1_updated"
obj2["field"] = "test_value_2_updated"
- self.hook.replace_many(collection, [obj1, obj2])
+ self.hook.replace_many("test_collection", [obj1, obj2],
mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == "test_value_1_updated"
- result_obj = collection.find_one(filter="2")
+ result_obj = mongo_collection.find_one(filter="2")
assert result_obj["field"] == "test_value_2_updated"
- def test_replace_many_with_upsert(self):
- collection = mongomock.MongoClient().db.collection
+ def test_replace_many_with_upsert(self, mongo_collection):
obj1 = {"_id": "1", "field": "test_value_1"}
obj2 = {"_id": "2", "field": "test_value_2"}
- self.hook.replace_many(collection, [obj1, obj2], upsert=True)
+ self.hook.replace_many("test_collection", [obj1, obj2], upsert=True,
mongo_db="test")
- result_obj = collection.find_one(filter="1")
+ result_obj = mongo_collection.find_one(filter="1")
assert result_obj["field"] == "test_value_1"
- result_obj = collection.find_one(filter="2")
+ result_obj = mongo_collection.find_one(filter="2")
assert result_obj["field"] == "test_value_2"
def test_create_uri_with_all_creds(self):
@@ -283,6 +274,7 @@ class TestMongoHook:
self.hook.connection.login = None
self.hook.connection.password = None
self.hook.connection.port = None
+ self.hook.connection.host = "mongo"
assert self.hook._create_uri() == "mongodb://mongo"
def test_create_uri_srv_true(self):
@@ -294,121 +286,118 @@ class TestMongoHook:
self.hook.connection.schema = "test_db"
assert self.hook._create_uri() ==
"mongodb+srv://test_user:test_password@test_host:1234/test_db"
- def test_delete_one(self):
- collection = mongomock.MongoClient().db.collection
+ def test_delete_one(self, mongo_collection):
obj = {"_id": "1"}
- collection.insert_one(obj)
+ mongo_collection.insert_one(obj)
- self.hook.delete_one(collection, {"_id": "1"})
+ self.hook.delete_one("test_collection", {"_id": "1"}, mongo_db="test")
- assert collection.count_documents({}) == 0
+ assert mongo_collection.count_documents({}) == 0
- def test_delete_many(self):
- collection = mongomock.MongoClient().db.collection
+ def test_delete_many(self, mongo_collection):
obj1 = {"_id": "1", "field": "value"}
obj2 = {"_id": "2", "field": "value"}
- collection.insert_many([obj1, obj2])
+ mongo_collection.insert_many([obj1, obj2])
- self.hook.delete_many(collection, {"field": "value"})
+ self.hook.delete_many("test_collection", {"field": "value"},
mongo_db="test")
- assert collection.count_documents({}) == 0
+ assert mongo_collection.count_documents({}) == 0
- def test_find_one(self):
- collection = mongomock.MongoClient().db.collection
+ def test_find_one(self, mongo_collection):
obj = {"test_find_one": "test_value"}
- collection.insert_one(obj)
+ mongo_collection.insert_one(obj)
- result_obj = self.hook.find(collection, {}, find_one=True)
+ result_obj = self.hook.find("test_collection", {}, find_one=True,
mongo_db="test")
result_obj = {result: result_obj[result] for result in result_obj}
assert obj == result_obj
- def test_find_many(self):
- collection = mongomock.MongoClient().db.collection
+ def test_find_many(self, mongo_collection):
objs = [{"_id": 1, "test_find_many_1": "test_value"}, {"_id": 2,
"test_find_many_2": "test_value"}]
- collection.insert_many(objs)
+ mongo_collection.insert_many(objs)
- result_objs = self.hook.find(mongo_collection=collection, query={},
projection={}, find_one=False)
+ result_objs = self.hook.find(
+ mongo_collection="test_collection", query={}, projection={},
find_one=False, mongo_db="test"
+ )
assert len(list(result_objs)) > 1
- def test_find_many_with_projection(self):
- collection = mongomock.MongoClient().db.collection
+ def test_find_many_with_projection(self, mongo_collection):
objs = [
{"_id": "1", "test_find_many_1": "test_value", "field_3": "a"},
{"_id": "2", "test_find_many_2": "test_value", "field_3": "b"},
]
- collection.insert_many(objs)
+ mongo_collection.insert_many(objs)
projection = {"_id": 0}
result_objs = self.hook.find(
- mongo_collection=collection, query={}, projection=projection,
find_one=False
+ mongo_collection="test_collection",
+ query={},
+ projection=projection,
+ find_one=False,
+ mongo_db="test",
)
assert "_id" not in result_objs[0]
- def test_aggregate(self):
- collection = mongomock.MongoClient().db.collection
+ def test_aggregate(self, mongo_collection):
objs = [
{"test_id": "1", "test_status": "success"},
{"test_id": "2", "test_status": "failure"},
{"test_id": "3", "test_status": "success"},
]
- collection.insert_many(objs)
+ mongo_collection.insert_many(objs)
aggregate_query = [{"$match": {"test_status": "success"}}]
- results = self.hook.aggregate(collection, aggregate_query)
+ results = self.hook.aggregate("test_collection", aggregate_query,
mongo_db="test")
assert len(list(results)) == 2
- def test_distinct(self):
- collection = mongomock.MongoClient().db.collection
+ def test_distinct(self, mongo_collection):
objs = [
{"test_id": "1", "test_status": "success"},
{"test_id": "2", "test_status": "failure"},
{"test_id": "3", "test_status": "success"},
]
- collection.insert_many(objs)
+ mongo_collection.insert_many(objs)
- results = self.hook.distinct(collection, "test_status")
+ results = self.hook.distinct("test_collection", "test_status",
mongo_db="test")
assert len(results) == 2
- def test_distinct_with_filter(self):
- collection = mongomock.MongoClient().db.collection
+ def test_distinct_with_filter(self, mongo_collection):
objs = [
{"test_id": "1", "test_status": "success"},
{"test_id": "2", "test_status": "failure"},
{"test_id": "3", "test_status": "success"},
]
- collection.insert_many(objs)
+ mongo_collection.insert_many(objs)
- results = self.hook.distinct(collection, "test_id", {"test_status":
"failure"})
+ results = self.hook.distinct(
+ "test_collection", "test_id", {"test_status": "failure"},
mongo_db="test"
+ )
assert len(results) == 1
- def test_create_standard_collection(self):
- mock_client = mongomock.MongoClient()
- self.hook.get_conn = lambda: mock_client
+ def test_create_standard_collection(self, container_mongo_client):
+ self.hook.get_conn = lambda: container_mongo_client
self.hook.connection.schema = "test_db"
collection =
self.hook.create_collection(mongo_collection="plain_collection")
assert collection.name == "plain_collection"
- assert "plain_collection" in
mock_client["test_db"].list_collection_names()
+ assert "plain_collection" in
container_mongo_client["test_db"].list_collection_names()
- def test_return_if_exists_true_returns_existing(self):
- mock_client = mongomock.MongoClient()
- self.hook.get_conn = lambda: mock_client
+ def test_return_if_exists_true_returns_existing(self,
container_mongo_client):
+ self.hook.get_conn = lambda: container_mongo_client
self.hook.connection.schema = "test_db"
first = self.hook.create_collection(mongo_collection="foo")
second = self.hook.create_collection(mongo_collection="foo",
return_if_exists=True)
assert first.full_name == second.full_name
- assert "foo" in mock_client["test_db"].list_collection_names()
+ assert "foo" in
container_mongo_client["test_db"].list_collection_names()
- def test_return_if_exists_false_raises(self):
- mock_client = mongomock.MongoClient()
- self.hook.get_conn = lambda: mock_client
+ def test_return_if_exists_false_raises(self, container_mongo_client):
+ self.hook.get_conn = lambda: container_mongo_client
self.hook.connection.schema = "test_db"
self.hook.create_collection(mongo_collection="bar")
diff --git a/scripts/ci/docker-compose/integration-mongo.yml
b/scripts/ci/docker-compose/integration-mongo.yml
index 0b35442fbc4..32a6b1c11f0 100644
--- a/scripts/ci/docker-compose/integration-mongo.yml
+++ b/scripts/ci/docker-compose/integration-mongo.yml
@@ -17,13 +17,13 @@
---
services:
mongo:
- image: mongo:3
+ image: mongo:latest
labels:
breeze.description: "Integration required for MongoDB hooks."
volumes:
- mongo-db-volume:/data/db
healthcheck:
- test: echo 'db.runCommand("ping").ok' | mongo localhost:27017/test
--quiet
+ test: echo 'db.runCommand("ping").ok' | mongosh localhost:27017/test
--quiet
interval: 5s
timeout: 30s
retries: 50
