PYTHON-1878 Add mongodb+srv URIs to Atlas Connectivity tests

.evergreen/config.yml

@@ -444,8 +444,29 @@ functions:
         silent: true
         working_dir: "src"
         script: |
-          # DO NOT ECHO WITH XTRACE (which PREPARE_SHELL does)
-          PYTHON_BINARY=${PYTHON_BINARY} ATLAS_REPL='${atlas_repl}' ATLAS_SHRD='${atlas_shrd}' ATLAS_FREE='${atlas_free}' ATLAS_TLS11='${atlas_tls11}' ATLAS_TLS12='${atlas_tls12}'  sh ${PROJECT_DIRECTORY}/.evergreen/run-atlas-tests.sh
+          cat <<EOT > prepare_atlas_connectivity.sh
+          export ATLAS_FREE='${atlas_free}'
+          export ATLAS_REPL='${atlas_repl}'
+          export ATLAS_SHRD='${atlas_shrd}'
+          export ATLAS_TLS11='${atlas_tls11}'
+          export ATLAS_TLS12='${atlas_tls12}'
+          export ATLAS_SRV_FREE='${atlas_srv_free}'
+          export ATLAS_SRV_REPL='${atlas_srv_repl}'
+          export ATLAS_SRV_SHRD='${atlas_srv_shrd}'
+          export ATLAS_SRV_TLS11='${atlas_srv_tls11}'
+          export ATLAS_SRV_TLS12='${atlas_srv_tls12}'
+          EOT
+    - command: shell.exec
+      type: test
+      params:
+        working_dir: "src"
+        script: |
+          # Disable xtrace (just in case it was accidentally set).
+          set +x
+          . ./prepare_atlas_connectivity.sh
+          rm -f ./prepare_atlas_connectivity.sh
+
+          PYTHON_BINARY=${PYTHON_BINARY} bash ${PROJECT_DIRECTORY}/.evergreen/run-atlas-tests.sh
 
   "add aws auth variables to file":
     - command: shell.exec

.evergreen/run-atlas-tests.sh

@@ -1,7 +1,8 @@
 #!/bin/bash
 
-# Don't trace to avoid secrets showing up in the logs
+# Exit on error and enable trace.
 set -o errexit
+set -o xtrace
 
 export JAVA_HOME=/opt/java/jdk8
 
@@ -15,27 +16,37 @@ if [ -z "$PYTHON_BINARY" ]; then
 fi
 
 IMPL=$(${PYTHON_BINARY} -c "import platform, sys; sys.stdout.write(platform.python_implementation())")
-if [ $IMPL = "Jython" -o $IMPL = "PyPy" ]; then
-    echo "Using Jython or PyPy"
+
+if [ $IMPL = "Jython" ]; then
+    # The venv created by createvirtualenv is incompatible with Jython
     $PYTHON_BINARY -m virtualenv --never-download --no-wheel atlastest
     . atlastest/bin/activate
-    trap "deactivate; rm -rf atlastest" EXIT HUP
-    pip install certifi
-    PYTHON=python
 else
-    IS_PRE_279=$(${PYTHON_BINARY} -c "import sys; sys.stdout.write('1' if sys.version_info < (2, 7, 9) else '0')")
+    # All other pythons work with createvirtualenv.
+    . .evergreen/utils.sh
+    createvirtualenv $PYTHON_BINARY atlastest
+fi
+trap "deactivate; rm -rf atlastest" EXIT HUP
+
+if [ $IMPL = "Jython" -o $IMPL = "PyPy" ]; then
+    echo "Using Jython or PyPy"
+    python -m pip install certifi
+else
+    IS_PRE_279=$(python -c "import sys; sys.stdout.write('1' if sys.version_info < (2, 7, 9) else '0')")
     if [ $IS_PRE_279 = "1" ]; then
         echo "Using a Pre-2.7.9 CPython"
-        $PYTHON_BINARY -m virtualenv --never-download --no-wheel atlastest
-        . atlastest/bin/activate
-        trap "deactivate; rm -rf atlastest" EXIT HUP
-        pip install pyopenssl>=17.2.0 service_identity>18.1.0
-        PYTHON=python
+        python -m pip install pyopenssl>=17.2.0 service_identity>18.1.0
     else
         echo "Using CPython 2.7.9+"
-        PYTHON=$PYTHON_BINARY
     fi
 fi
 
-echo "Running tests"
-$PYTHON test/atlas/test_connection.py
+echo "Running tests without dnspython"
+python test/atlas/test_connection.py
+
+# dnspython is incompatible with Jython so don't test that combination.
+if [ $IMPL != "Jython" ]; then
+    python -m pip install dnspython
+    echo "Running tests with dnspython"
+    MUST_TEST_SRV="1" python test/atlas/test_connection.py
+fi

test/atlas/test_connection.py

@@ -18,20 +18,41 @@
 import sys
 import unittest
 
+from collections import defaultdict
+
 sys.path[0:0] = [""]
 
 import pymongo
 from pymongo.ssl_support import HAS_SNI
 
-
-_REPL = os.environ.get("ATLAS_REPL")
-_SHRD = os.environ.get("ATLAS_SHRD")
-_FREE = os.environ.get("ATLAS_FREE")
-_TLS11 = os.environ.get("ATLAS_TLS11")
-_TLS12 = os.environ.get("ATLAS_TLS12")
-
-
-def _connect(uri):
+try:
+    import dns
+    HAS_DNS = True
+except ImportError:
+    HAS_DNS = False
+
+
+URIS = {
+    "ATLAS_REPL": os.environ.get("ATLAS_REPL"),
+    "ATLAS_SHRD": os.environ.get("ATLAS_SHRD"),
+    "ATLAS_FREE": os.environ.get("ATLAS_FREE"),
+    "ATLAS_TLS11": os.environ.get("ATLAS_TLS11"),
+    "ATLAS_TLS12": os.environ.get("ATLAS_TLS12"),
+    "ATLAS_SRV_REPL": os.environ.get("ATLAS_SRV_REPL"),
+    "ATLAS_SRV_SHRD": os.environ.get("ATLAS_SRV_SHRD"),
+    "ATLAS_SRV_FREE": os.environ.get("ATLAS_SRV_FREE"),
+    "ATLAS_SRV_TLS11": os.environ.get("ATLAS_SRV_TLS11"),
+    "ATLAS_SRV_TLS12": os.environ.get("ATLAS_SRV_TLS12"),
+}
+
+# Set this variable to true to run the SRV tests even when dnspython is not
+# installed.
+MUST_TEST_SRV = os.environ.get("MUST_TEST_SRV")
+
+
+def connect(uri):
+    if not uri:
+        raise Exception("Must set env variable to test.")
     client = pymongo.MongoClient(uri)
     # No TLS error
     client.admin.command('ismaster')
@@ -40,29 +61,58 @@ def _connect(uri):
 
 
 class TestAtlasConnect(unittest.TestCase):
-
-    @classmethod
-    def setUpClass(cls):
-        if not all([_REPL, _SHRD, _FREE]):
-            raise Exception(
-                "Must set ATLAS_REPL/SHRD/FREE env variables to test.")
+    @unittest.skipUnless(HAS_SNI, 'Free tier requires SNI support')
+    def test_free_tier(self):
+        connect(URIS['ATLAS_FREE'])
 
     def test_replica_set(self):
-        _connect(_REPL)
+        connect(URIS['ATLAS_REPL'])
 
     def test_sharded_cluster(self):
-        _connect(_SHRD)
-
-    def test_free_tier(self):
-        if not HAS_SNI:
-            raise unittest.SkipTest("Free tier requires SNI support.")
-        _connect(_FREE)
+        connect(URIS['ATLAS_SHRD'])
 
     def test_tls_11(self):
-        _connect(_TLS11)
+        connect(URIS['ATLAS_TLS11'])
 
     def test_tls_12(self):
-        _connect(_TLS12)
+        connect(URIS['ATLAS_TLS12'])
+
+    def connect_srv(self, uri):
+        connect(uri)
+        self.assertIn('mongodb+srv://', uri)
+
+    @unittest.skipUnless(HAS_SNI, 'Free tier requires SNI support')
+    @unittest.skipUnless(HAS_DNS or MUST_TEST_SRV, 'SRV requires dnspython')
+    def test_srv_free_tier(self):
+        self.connect_srv(URIS['ATLAS_SRV_FREE'])
+
+    @unittest.skipUnless(HAS_DNS or MUST_TEST_SRV, 'SRV requires dnspython')
+    def test_srv_replica_set(self):
+        self.connect_srv(URIS['ATLAS_SRV_REPL'])
+
+    @unittest.skipUnless(HAS_DNS or MUST_TEST_SRV, 'SRV requires dnspython')
+    def test_srv_sharded_cluster(self):
+        self.connect_srv(URIS['ATLAS_SRV_SHRD'])
+
+    @unittest.skipUnless(HAS_DNS or MUST_TEST_SRV, 'SRV requires dnspython')
+    def test_srv_tls_11(self):
+        self.connect_srv(URIS['ATLAS_SRV_TLS11'])
+
+    @unittest.skipUnless(HAS_DNS or MUST_TEST_SRV, 'SRV requires dnspython')
+    def test_srv_tls_12(self):
+        self.connect_srv(URIS['ATLAS_SRV_TLS12'])
+
+    def test_uniqueness(self):
+        """Ensure that we don't accidentally duplicate the test URIs."""
+        uri_to_names = defaultdict(list)
+        for name, uri in URIS.items():
+            if uri:
+                uri_to_names[uri].append(name)
+        duplicates = [names for names in uri_to_names.values()
+                      if len(names) > 1]
+        self.assertFalse(duplicates, 'Error: the following env variables have '
+                                     'duplicate values: %s' % (duplicates,))
+
 
 
 if __name__ == '__main__':