libAtoms · stenczelt · Jun 14, 2024 · Jun 13, 2024 · Jun 13, 2024 · Jun 13, 2024
diff --git a/abcd/__init__.py b/abcd/__init__.py
@@ -39,7 +39,12 @@ def from_url(cls, url, **kwargs):
             from abcd.backends.atoms_pymongo import MongoDatabase
 
             return MongoDatabase(db_name=db, **conn_settings, **kwargs)
+        elif r.scheme == "mongodb+srv":
+            db = r.path.split("/")[1] if r.path else None
+            db = db if db else "abcd"
+            from abcd.backends.atoms_pymongo import MongoDatabase
 
+            return MongoDatabase(db_name=db, host=r.geturl(), uri_mode=True, **kwargs)
         elif r.scheme == "http" or r.scheme == "https":
             raise NotImplementedError("http not yet supported! soon...")
         elif r.scheme == "ssh":

diff --git a/abcd/backends/atoms_pymongo.py b/abcd/backends/atoms_pymongo.py
@@ -164,6 +164,7 @@ def __init__(
         username=None,
         password=None,
         authSource="admin",
+        uri_mode=False,
         **kwargs
     ):
         super().__init__()
@@ -181,13 +182,16 @@ def __init__(
             )
         )
 
-        self.client = MongoClient(
-            host=host,
-            port=port,
-            username=username,
-            password=password,
-            authSource=authSource,
-        )
+        if uri_mode:
+            self.client = MongoClient(host=host, authSource=authSource)
+        else:
+            self.client = MongoClient(
+                host=host,
+                port=port,
+                username=username,
+                password=password,
+                authSource=authSource,
+            )
 
         try:
             info = self.client.server_info()  # Forces a call.

diff --git a/pyproject.toml b/pyproject.toml
@@ -23,6 +23,7 @@ lark = "^1.1.9"
 mongomock = "^4.1.2"
 pytest = "^8.2.2"
 pytest-cov = "^5.0.0"
+pytest-mock = "^3.14.0"
 
 [tool.poetry.extras]
 tests = ["mongomock", "pytest", "pytest-cov"]

diff --git a/tests/test_mongodb_srv.py b/tests/test_mongodb_srv.py
@@ -0,0 +1,29 @@
+"""Tests for supporting `mongodb+srv://` URIs"""
+
+from pytest import fixture
+
+from abcd import ABCD
+
+
+@fixture
+def mongo_srv(mocker):
+    # mongomock does not pick up what we need, so this is a bespoke mocker
+    mock_client = mocker.MagicMock()
+    mocker.patch("abcd.backends.atoms_pymongo.MongoClient", mock_client)
+    return mock_client
+
+
+def test_init_mongodb_srv(mongo_srv):
+    # client can be created with a mongodb+srv:// URI
+
+    # apparently mongomock breaks if this import is outside
+    from abcd.backends.atoms_pymongo import MongoDatabase
+
+    # regression test
+    uri = "mongodb+srv://user:pass@democluster.randomstr.mongodb.net/?key=value"
+
+    # create the client
+    abcd = ABCD.from_url(uri)
+
+    assert isinstance(abcd, MongoDatabase)
+    mongo_srv.assert_called_once_with(host=uri, authSource="admin")