Fix test cases which can run on remote server

Author: passren

pymongosql/sqlalchemy_mongodb/__init__.py

@@ -128,10 +128,20 @@ def register_dialect():
     try:
         from sqlalchemy.dialects import registry
 
-        # Register for standard MongoDB URLs only
+        # Register for standard MongoDB URLs
         registry.register("mongodb", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
-        # Note: mongodb+srv is handled by converting to mongodb in create_connect_args
-        # SQLAlchemy doesn't support the + character in scheme names directly
+
+        # Try to register both SRV forms so SQLAlchemy can resolve SRV-style URLs
+        # (either 'mongodb+srv' or the dotted 'mongodb.srv' plugin name).
+        # Some SQLAlchemy versions accept '+' in scheme names; others import
+        # the dotted plugin name. Attempt both registrations in one block.
+        try:
+            registry.register("mongodb+srv", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
+            registry.register("mongodb.srv", "pymongosql.sqlalchemy_mongodb.sqlalchemy_dialect", "PyMongoSQLDialect")
+        except Exception:
+            # If registration fails we fall back to handling SRV URIs in
+            # create_engine_from_mongodb_uri by converting 'mongodb+srv' to 'mongodb'.
+            pass
 
         return True
     except ImportError:

tests/test_sqlalchemy_integration.py

@@ -9,8 +9,12 @@
 4. Validating object creation from query results
 """
 
+import os
+
 import pytest
 
+from tests.conftest import TEST_DB, TEST_URI
+
 # SQLAlchemy version compatibility
 try:
     import sqlalchemy
@@ -107,8 +111,32 @@ def __repr__(self):
 # Pytest fixtures
 @pytest.fixture
 def sqlalchemy_engine():
-    """Provide a SQLAlchemy engine connected to MongoDB."""
-    engine = create_engine("mongodb://testuser:testpass@localhost:27017/test_db")
+    """Provide a SQLAlchemy engine connected to MongoDB. The URI is taken from environment variables
+    (PYMONGOSQL_TEST_URI or MONGODB_URI) or falls back to a sensible local default.
+    """
+    uri = os.environ.get("PYMONGOSQL_TEST_URI") or os.environ.get("MONGODB_URI") or TEST_URI
+    db = os.environ.get("PYMONGOSQL_TEST_DB") or TEST_DB
+
+    def _ensure_uri_has_db(uri_value: str, database: str) -> str:
+        if not database:
+            return uri_value
+        idx = uri_value.find("://")
+        if idx == -1:
+            return uri_value
+        rest = uri_value[idx + 3 :]
+        if "/" in rest:
+            after = rest.split("/", 1)[1]
+            if after == "" or after.startswith("?"):
+                return uri_value.rstrip("/") + "/" + database
+            return uri_value
+        return uri_value.rstrip("/") + "/" + database
+
+    if uri:
+        uri_to_use = _ensure_uri_has_db(uri, db)
+    else:
+        uri_to_use = "mongodb://testuser:testpass@localhost:27017/test_db"
+
+    engine = create_engine(uri_to_use)
     yield engine