PYTHON-5206 Convert Atlas Connect and Enterprise Auth tests to use new test scripts (#2201)

Co-authored-by: Noah Stapp <noah@noahstapp.com>

Author: blink1073

.evergreen/config.yml

@@ -267,28 +267,6 @@ functions:
         binary: bash
         args: [.evergreen/just.sh, run-tests]
 
-  "run enterprise auth tests":
-    - command: subprocess.exec
-      type: test
-      params:
-        binary: bash
-        working_dir: "src"
-        include_expansions_in_env: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "PYTHON_BINARY"]
-        args:
-          - .evergreen/scripts/run-with-env.sh
-          - .evergreen/scripts/run-enterprise-auth-tests.sh
-
-  "run atlas tests":
-    - command: subprocess.exec
-      type: test
-      params:
-        binary: bash
-        include_expansions_in_env: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "PYTHON_BINARY"]
-        working_dir: "src"
-        args:
-          - .evergreen/scripts/run-with-env.sh
-          - .evergreen/scripts/run-atlas-tests.sh
-
   "cleanup":
     - command: subprocess.exec
       params:
@@ -422,13 +400,6 @@ tasks:
         - func: "run server"
         - func: "run doctests"
 
-    - name: "test-enterprise-auth"
-      tags: ["enterprise-auth"]
-      commands:
-        - func: "run server"
-        - func: "assume ec2 role"
-        - func: "run enterprise auth tests"
-
     - name: "test-search-index-helpers"
       commands:
         - func: "run server"
@@ -488,12 +459,6 @@ tasks:
             TOPOLOGY: "replica_set"
         - func: "run tests"
 
-    - name: "atlas-connect"
-      tags: ["atlas-connect"]
-      commands:
-        - func: "assume ec2 role"
-        - func: "run atlas tests"
-
     - name: atlas-data-lake-tests
       commands:
         - func: "bootstrap data lake"

.evergreen/generated_configs/tasks.yml

@@ -1,4 +1,13 @@
 tasks:
+  # Atlas connect tests
+  - name: test-atlas-connect
+    commands:
+      - func: assume ec2 role
+      - func: run tests
+        vars:
+          TEST_NAME: atlas_connect
+    tags: [atlas_connect]
+
   # Aws tests
   - name: test-auth-aws-4.4-regular
     commands:
@@ -680,6 +689,20 @@ tasks:
           AWS_ROLE_SESSION_NAME: test
     tags: [auth-aws, auth-aws-web-identity]
 
+  # Enterprise auth tests
+  - name: test-enterprise-auth
+    commands:
+      - func: run server
+        vars:
+          TEST_NAME: enterprise_auth
+          AUTH: auth
+      - func: assume ec2 role
+      - func: run tests
+        vars:
+          TEST_NAME: enterprise_auth
+          AUTH: auth
+    tags: [enterprise_auth]
+
   # Kms tests
   - name: test-gcpkms
     commands:

.evergreen/generated_configs/variants.yml

@@ -49,15 +49,15 @@ buildvariants:
   # Atlas connect tests
   - name: atlas-connect-rhel8-python3.9
     tasks:
-      - name: atlas-connect
+      - name: .atlas_connect
     display_name: Atlas connect RHEL8 Python3.9
     run_on:
       - rhel87-small
     expansions:
       PYTHON_BINARY: /opt/python/3.9/bin/python3
   - name: atlas-connect-rhel8-python3.13
     tasks:
-      - name: atlas-connect
+      - name: .atlas_connect
     display_name: Atlas connect RHEL8 Python3.13
     run_on:
       - rhel87-small
@@ -510,59 +510,53 @@ buildvariants:
     tags: [encryption_tag]
 
   # Enterprise auth tests
-  - name: auth-enterprise-macos-python3.9-auth
+  - name: auth-enterprise-macos-python3.9
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise macOS Python3.9 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise macOS Python3.9
     run_on:
       - macos-14
     expansions:
-      AUTH: auth
       PYTHON_BINARY: /Library/Frameworks/Python.Framework/Versions/3.9/bin/python3
-  - name: auth-enterprise-rhel8-python3.10-auth
+  - name: auth-enterprise-rhel8-python3.10
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise RHEL8 Python3.10 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise RHEL8 Python3.10
     run_on:
       - rhel87-small
     expansions:
-      AUTH: auth
       PYTHON_BINARY: /opt/python/3.10/bin/python3
-  - name: auth-enterprise-rhel8-python3.11-auth
+  - name: auth-enterprise-rhel8-python3.11
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise RHEL8 Python3.11 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise RHEL8 Python3.11
     run_on:
       - rhel87-small
     expansions:
-      AUTH: auth
       PYTHON_BINARY: /opt/python/3.11/bin/python3
-  - name: auth-enterprise-rhel8-python3.12-auth
+  - name: auth-enterprise-rhel8-python3.12
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise RHEL8 Python3.12 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise RHEL8 Python3.12
     run_on:
       - rhel87-small
     expansions:
-      AUTH: auth
       PYTHON_BINARY: /opt/python/3.12/bin/python3
-  - name: auth-enterprise-win64-python3.13-auth
+  - name: auth-enterprise-win64-python3.13
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise Win64 Python3.13 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise Win64 Python3.13
     run_on:
       - windows-64-vsMulti-small
     expansions:
-      AUTH: auth
       PYTHON_BINARY: C:/python/Python313/python.exe
-  - name: auth-enterprise-rhel8-pypy3.10-auth
+  - name: auth-enterprise-rhel8-pypy3.10
     tasks:
-      - name: test-enterprise-auth
-    display_name: Auth Enterprise RHEL8 PyPy3.10 Auth
+      - name: .enterprise_auth
+    display_name: Auth Enterprise RHEL8 PyPy3.10
     run_on:
       - rhel87-small
     expansions:
-      AUTH: auth
       PYTHON_BINARY: /opt/python/pypy3.10/bin/python3
 
   # Free threaded tests

.evergreen/scripts/generate_config.py

@@ -464,7 +464,6 @@ def create_compression_variants():
 
 
 def create_enterprise_auth_variants():
-    expansions = dict(AUTH="auth")
     variants = []
 
     # All python versions across platforms.
@@ -475,10 +474,8 @@ def create_enterprise_auth_variants():
             host = HOSTS["win64"]
         else:
             host = DEFAULT_HOST
-        display_name = get_display_name("Auth Enterprise", host, python=python, **expansions)
-        variant = create_variant(
-            ["test-enterprise-auth"], display_name, host=host, python=python, expansions=expansions
-        )
+        display_name = get_display_name("Auth Enterprise", host, python=python)
+        variant = create_variant([".enterprise_auth"], display_name, host=host, python=python)
         variants.append(variant)
 
     return variants
@@ -721,7 +718,7 @@ def create_atlas_connect_variants():
     host = DEFAULT_HOST
     return [
         create_variant(
-            ["atlas-connect"],
+            [".atlas_connect"],
             get_display_name("Atlas connect", host, python=python),
             python=python,
             host=host,
@@ -913,6 +910,25 @@ def _create_ocsp_task(algo, variant, server_type, base_task_name):
     return EvgTask(name=task_name, tags=tags, commands=commands)
 
 
+def create_atlas_connect_tasks():
+    vars = dict(TEST_NAME="atlas_connect")
+    assume_func = FunctionCall(func="assume ec2 role")
+    test_func = FunctionCall(func="run tests", vars=vars)
+    task_name = "test-atlas-connect"
+    tags = ["atlas_connect"]
+    return [EvgTask(name=task_name, tags=tags, commands=[assume_func, test_func])]
+
+
+def create_enterprise_auth_tasks():
+    vars = dict(TEST_NAME="enterprise_auth", AUTH="auth")
+    server_func = FunctionCall(func="run server", vars=vars)
+    assume_func = FunctionCall(func="assume ec2 role")
+    test_func = FunctionCall(func="run tests", vars=vars)
+    task_name = "test-enterprise-auth"
+    tags = ["enterprise_auth"]
+    return [EvgTask(name=task_name, tags=tags, commands=[server_func, assume_func, test_func])]
+
+
 def create_ocsp_tasks():
     tasks = []
     tests = [

.evergreen/scripts/run-atlas-tests.sh

@@ -112,6 +112,10 @@ def setup_libmongocrypt():
         run_command("chmod +x libmongocrypt/nocrypto/bin/mongocrypt.dll")
 
 
+def get_secrets(name: str) -> None:
+    run_command(f"bash {DRIVERS_TOOLS}/.evergreen/secrets_handling/setup-secrets.sh {name}")
+
+
 def handle_test_env() -> None:
     opts, _ = get_test_options("Set up the test environment and services.")
     test_name = opts.test_name
@@ -203,6 +207,7 @@ def handle_test_env() -> None:
         write_env("PYMONGO_DISABLE_TEST_COMMANDS", "1")
 
     if test_name == "enterprise_auth":
+        get_secrets("drivers/enterprise_auth")
         config = read_env(f"{ROOT}/secrets-export.sh")
         if PLATFORM == "windows":
             LOGGER.info("Setting GSSAPI_PASS")
@@ -346,6 +351,11 @@ def handle_test_env() -> None:
         else:
             run_command(f"bash {auth_aws_dir}/setup-secrets.sh")
 
+    if test_name == "atlas_connect":
+        get_secrets("drivers/atlas_connect")
+        # We do not want the default client_context to be initialized.
+        write_env("DISABLE_CONTEXT")
+
     if test_name == "perf":
         # PYTHON-4769 Run perf_test.py directly otherwise pytest's test collection negatively
         # affects the benchmark results.

.evergreen/scripts/run-direct-tests.sh

@@ -30,7 +30,7 @@ class Distro:
 
 # Map the test name to a test suite.
 TEST_SUITE_MAP = {
-    "atlas": "atlas",
+    "atlas_connect": "atlas_connect",
     "auth_aws": "auth_aws",
     "auth_oidc": "auth_oidc",
     "data_lake": "data_lake",

.evergreen/scripts/run-enterprise-auth-tests.sh

@@ -262,6 +262,19 @@ For KMS tests that run remotely and are expected to pass, in this case using `gc
 - Run `just setup-tests kms gcp`.
 - Run `just run-tests`.
 
+### Enterprise Auth tests
+
+Note: these tests can only be run from an Evergreen host.
+
+- Run `just run-server enterprise_auth`.
+- Run `just setup-tests enterprise_auth`.
+- Run `just run-tests`.
+
+### Atlas Connect tests
+
+- Run `just setup-tests atlas_connect`.
+- Run `just run-tests`.
+
 ### OCSP tests
 
   - Export the orchestration file, e.g. `export ORCHESTRATION_FILE=rsa-basic-tls-ocsp-disableStapling.json`.