PYTHON-5213 Convert AWS Lambda and Search Index tests to use new test scripts (#2205)

Author: blink1073

.evergreen/config.yml

@@ -268,27 +268,7 @@ functions:
     - command: ec2.assume_role
       params:
         role_arn: ${aws_test_secrets_role}
-
-  "setup atlas":
-    - command: subprocess.exec
-      params:
-        binary: bash
-        include_expansions_in_env: ["task_id", "execution"]
-        env:
-          MONGODB_VERSION: "7.0"
-          LAMBDA_STACK_NAME: dbx-python-lambda
-        args:
-          - ${DRIVERS_TOOLS}/.evergreen/atlas/setup-atlas-cluster.sh
-    - command: expansions.update
-      params:
-        file: atlas-expansion.yml
-
-  "teardown atlas":
-    - command: subprocess.exec
-      params:
-        binary: bash
-        args:
-          - ${DRIVERS_TOOLS}/.evergreen/atlas/teardown-atlas-cluster.sh
+        duration_seconds: 3600
 
   "attach benchmark test results":
     - command: attach.results
@@ -314,31 +294,6 @@ post:
   - func: "upload test results"
   - func: "cleanup"
 
-task_groups:
-  - name: test_aws_lambda_task_group
-    setup_group:
-      - func: fetch source
-      - func: setup system
-      - func: setup atlas
-    teardown_task:
-      - func: teardown atlas
-    setup_group_can_fail_task: true
-    setup_group_timeout_secs: 1800
-    tasks:
-      - test-aws-lambda-deployed
-
-  - name: test_atlas_task_group_search_indexes
-    setup_group:
-      - func: fetch source
-      - func: setup system
-      - func: setup atlas
-    teardown_task:
-      - func: teardown atlas
-    setup_group_can_fail_task: true
-    setup_group_timeout_secs: 1800
-    tasks:
-      - test-search-index-helpers
-
 tasks:
     # Wildcard task. Do you need to find out what tools are available and where?
     # Throw it here, and execute this task on all buildvariants
@@ -365,17 +320,6 @@ tasks:
         - func: "run server"
         - func: "run doctests"
 
-    - name: "test-search-index-helpers"
-      commands:
-        - func: "run server"
-          vars:
-            VERSION: "6.0"
-            TOPOLOGY: "replica_set"
-        - func: "run tests"
-          vars:
-            TEST_NAME: index_management
-            AUTH: "auth"
-
     - name: "no-server"
       tags: ["no-server"]
       commands:
@@ -390,22 +334,6 @@ tasks:
             TOPOLOGY: "replica_set"
         - func: "run tests"
 
-    - name: "test-aws-lambda-deployed"
-      commands:
-        - command: ec2.assume_role
-          params:
-            role_arn: ${LAMBDA_AWS_ROLE_ARN}
-            duration_seconds: 3600
-        - command: subprocess.exec
-          params:
-            working_dir: src
-            binary: bash
-            add_expansions_to_env: true
-            args:
-              - .evergreen/run-deployed-lambda-aws-tests.sh
-            env:
-              TEST_LAMBDA_DIRECTORY: ${PROJECT_DIRECTORY}/test/lambda
-
 # }}}
     - name: "coverage-report"
       tags: ["coverage"]
@@ -482,12 +410,6 @@ buildvariants:
       batchtime: 10080  # 7 days
     - name: test-azurekms-fail
 
-- name: rhel8-test-lambda
-  display_name: FaaS Lambda
-  run_on: rhel87-small
-  tasks:
-    - name: test_aws_lambda_task_group
-
 - name: rhel8-import-time
   display_name: Import Time
   run_on: rhel87-small

.evergreen/generated_configs/tasks.yml

@@ -23,6 +23,15 @@ tasks:
           TEST_NAME: data_lake
     tags: [atlas_data_lake]
 
+  # Aws lambda tests
+  - name: test-aws-lambda-deployed
+    commands:
+      - func: assume ec2 role
+      - func: run tests
+        vars:
+          TEST_NAME: aws_lambda
+    tags: [aws_lambda]
+
   # Aws tests
   - name: test-auth-aws-4.4-regular
     commands:
@@ -1220,6 +1229,18 @@ tasks:
       - func: send dashboard data
     tags: [perf]
 
+  # Search index tests
+  - name: test-search-index-helpers
+    commands:
+      - func: assume ec2 role
+      - func: run server
+        vars:
+          TEST_NAME: search_index
+      - func: run tests
+        vars:
+          TEST_NAME: search_index
+    tags: [search_index]
+
   # Server tests
   - name: test-4.0-standalone-auth-ssl-sync
     commands:

.evergreen/generated_configs/variants.yml

@@ -132,6 +132,14 @@ buildvariants:
     expansions:
       PYTHON_BINARY: /Library/Frameworks/Python.Framework/Versions/3.13/bin/python3
 
+  # Aws lambda tests
+  - name: faas-lambda
+    tasks:
+      - name: .aws_lambda
+    display_name: FaaS Lambda
+    run_on:
+      - rhel87-small
+
   # Compression tests
   - name: compression-snappy-rhel8-python3.9-no-c
     tasks:
@@ -986,7 +994,7 @@ buildvariants:
   # Search index tests
   - name: search-index-helpers-rhel8-python3.9
     tasks:
-      - name: test_atlas_task_group_search_indexes
+      - name: .search_index
     display_name: Search Index Helpers RHEL8 Python3.9
     run_on:
       - rhel87-small

.evergreen/run-deployed-lambda-aws-tests.sh

@@ -672,7 +672,7 @@ def create_search_index_variants():
     python = CPYTHONS[0]
     return [
         create_variant(
-            ["test_atlas_task_group_search_indexes"],
+            [".search_index"],
             get_display_name("Search Index Helpers", host, python=python),
             python=python,
             host=host,
@@ -779,6 +779,11 @@ def create_alternative_hosts_variants():
     return variants
 
 
+def create_aws_lambda_variants():
+    host = HOSTS["rhel8"]
+    return [create_variant([".aws_lambda"], display_name="FaaS Lambda", host=host)]
+
+
 ##############
 # Tasks
 ##############
@@ -927,6 +932,27 @@ def _create_ocsp_task(algo, variant, server_type, base_task_name):
     return EvgTask(name=task_name, tags=tags, commands=commands)
 
 
+def create_aws_lambda_tasks():
+    assume_func = FunctionCall(func="assume ec2 role")
+    vars = dict(TEST_NAME="aws_lambda")
+    test_func = FunctionCall(func="run tests", vars=vars)
+    task_name = "test-aws-lambda-deployed"
+    tags = ["aws_lambda"]
+    commands = [assume_func, test_func]
+    return [EvgTask(name=task_name, tags=tags, commands=commands)]
+
+
+def create_search_index_tasks():
+    assume_func = FunctionCall(func="assume ec2 role")
+    server_func = FunctionCall(func="run server", vars=dict(TEST_NAME="search_index"))
+    vars = dict(TEST_NAME="search_index")
+    test_func = FunctionCall(func="run tests", vars=vars)
+    task_name = "test-search-index-helpers"
+    tags = ["search_index"]
+    commands = [assume_func, server_func, test_func]
+    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")

.evergreen/scripts/generate_config.py

@@ -22,14 +22,6 @@ def start_server():
     if "VERSION" in os.environ:
         os.environ["MONGODB_VERSION"] = os.environ["VERSION"]
 
-    if opts.auth:
-        extra_opts.append("--auth")
-
-    if opts.verbose:
-        extra_opts.append("-v")
-    elif opts.quiet:
-        extra_opts.append("-q")
-
     if test_name == "auth_aws":
         set_env("AUTH_AWS")
 
@@ -51,6 +43,10 @@ def start_server():
             if not found:
                 raise ValueError("Please provide an orchestration file")
 
+    elif test_name == "search_index":
+        os.environ["TOPOLOGY"] = "replica_set"
+        os.environ["MONGODB_VERSION"] = "7.0"
+
     if not os.environ.get("TEST_CRYPT_SHARED"):
         set_env("SKIP_CRYPT_SHARED")
 
@@ -62,6 +58,14 @@ def start_server():
             set_env("TLS_PEM_KEY_FILE", certs / "server.pem")
             set_env("TLS_CA_FILE", certs / "ca.pem")
 
+    if opts.auth:
+        extra_opts.append("--auth")
+
+    if opts.verbose:
+        extra_opts.append("-v")
+    elif opts.quiet:
+        extra_opts.append("-q")
+
     cmd = ["bash", f"{DRIVERS_TOOLS}/.evergreen/run-orchestration.sh", *extra_opts]
     run_command(cmd, cwd=DRIVERS_TOOLS)
 

.evergreen/scripts/run_server.py

@@ -4,8 +4,11 @@
 import logging
 import os
 import platform
+import shutil
 import sys
 from datetime import datetime
+from pathlib import Path
+from shutil import which
 
 import pytest
 from utils import DRIVERS_TOOLS, LOGGER, ROOT, run_command
@@ -81,6 +84,42 @@ def handle_pymongocrypt() -> None:
     LOGGER.info(f"libmongocrypt version: {pymongocrypt.libmongocrypt_version()})")
 
 
+def handle_aws_lambda() -> None:
+    env = os.environ.copy()
+    target_dir = ROOT / "test/lambda"
+    env["TEST_LAMBDA_DIRECTORY"] = str(target_dir)
+    env.setdefault("AWS_REGION", "us-east-1")
+    dirs = ["pymongo", "gridfs", "bson"]
+    # Store the original .so files.
+    before_sos = []
+    for dname in dirs:
+        before_sos.extend(f"{f.parent.name}/{f.name}" for f in (ROOT / dname).glob("*.so"))
+    # Build the c extensions.
+    docker = which("docker") or which("podman")
+    if not docker:
+        raise ValueError("Could not find docker!")
+    image = "quay.io/pypa/manylinux2014_x86_64:latest"
+    run_command(
+        f'{docker} run --rm -v "{ROOT}:/src" --platform linux/amd64 {image} /src/test/lambda/build_internal.sh'
+    )
+    for dname in dirs:
+        target = ROOT / "test/lambda/mongodb" / dname
+        shutil.rmtree(target, ignore_errors=True)
+        shutil.copytree(ROOT / dname, target)
+    # Remove the original so files from the lambda directory.
+    for so_path in before_sos:
+        (ROOT / "test/lambda/mongodb" / so_path).unlink()
+    # Remove the new so files from the ROOT directory.
+    for dname in dirs:
+        so_paths = [f"{f.parent.name}/{f.name}" for f in (ROOT / dname).glob("*.so")]
+        for so_path in list(so_paths):
+            if so_path not in before_sos:
+                Path(so_path).unlink()
+
+    script_name = "run-deployed-lambda-aws-tests.sh"
+    run_command(f"bash {DRIVERS_TOOLS}/.evergreen/aws_lambda/{script_name}", env=env)
+
+
 def run() -> None:
     # Handle green framework first so they can patch modules.
     if GREEN_FRAMEWORK:
@@ -129,6 +168,11 @@ def run() -> None:
         test_oidc_send_to_remote(SUB_TEST_NAME)
         return
 
+    # Run deployed aws lambda tests.
+    if TEST_NAME == "aws_lambda":
+        handle_aws_lambda()
+        return
+
     if os.environ.get("DEBUG_LOG"):
         TEST_ARGS.extend(f"-o log_cli_level={logging.DEBUG} -o log_cli=1".split())
 

.evergreen/scripts/run_tests.py

@@ -175,6 +175,28 @@ def handle_test_env() -> None:
         if not config:
             AUTH = "noauth"
 
+    if test_name in ["aws_lambda", "search_index"]:
+        env = os.environ.copy()
+        env["MONGODB_VERSION"] = "7.0"
+        env["LAMBDA_STACK_NAME"] = "dbx-python-lambda"
+        write_env("LAMBDA_STACK_NAME", env["LAMBDA_STACK_NAME"])
+        run_command(
+            f"bash {DRIVERS_TOOLS}/.evergreen/atlas/setup-atlas-cluster.sh",
+            env=env,
+            cwd=DRIVERS_TOOLS,
+        )
+
+    if test_name == "search_index":
+        AUTH = "auth"
+
+    if test_name == "aws_lambda":
+        UV_ARGS.append("--with pip")
+        # Store AWS creds if they were given.
+        if "AWS_ACCESS_KEY_ID" in os.environ:
+            for key in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]:
+                if key in os.environ:
+                    write_env(key, os.environ[key])
+
     if test_name == "data_lake":
         # Stop any running mongo-orchestration which might be using the port.
         run_command(f"bash {DRIVERS_TOOLS}/.evergreen/stop-orchestration.sh")
@@ -197,7 +219,7 @@ def handle_test_env() -> None:
         elif test_name == "auth_oidc":
             DB_USER = config["OIDC_ADMIN_USER"]
             DB_PASSWORD = config["OIDC_ADMIN_PWD"]
-        elif test_name == "index_management":
+        elif test_name == "search_index":
             config = read_env(f"{DRIVERS_TOOLS}/.evergreen/atlas/secrets-export.sh")
             DB_USER = config["DRIVERS_ATLAS_LAMBDA_USER"]
             DB_PASSWORD = config["DRIVERS_ATLAS_LAMBDA_PASSWORD"]

.evergreen/scripts/setup_tests.py

@@ -40,6 +40,10 @@
 elif TEST_NAME == "serverless":
     run_command(f"bash {DRIVERS_TOOLS}/.evergreen/serverless/teardown.sh")
 
+# Tear down atlas cluster if applicable.
+if TEST_NAME in ["aws_lambda", "search_index"]:
+    run_command(f"bash {DRIVERS_TOOLS}/.evergreen/atlas/teardown-atlas-cluster.sh")
+
 # Tear down auth_aws if applicable.
 # We do not run web-identity hosts on macos, because the hosts lack permissions,
 # so there is no reason to run the teardown, which would error with a 401.