From 5e055eea0f09b0a2856063af80af810134d23849 Mon Sep 17 00:00:00 2001 From: Steven Silvester Date: Thu, 13 Mar 2025 12:00:56 -0500 Subject: [PATCH 1/2] PYTHON-5206 Convert Atlas Connect and Enterprise Auth tests to use new test scripts (#2201) Co-authored-by: Noah Stapp --- .evergreen/config.yml | 35 -------------- .evergreen/generated_configs/tasks.yml | 23 ++++++++++ .evergreen/generated_configs/variants.yml | 46 ++++++++----------- .evergreen/scripts/generate_config.py | 28 ++++++++--- .evergreen/scripts/run-atlas-tests.sh | 8 ---- .evergreen/scripts/run-direct-tests.sh | 10 ---- .../scripts/run-enterprise-auth-tests.sh | 9 ---- .evergreen/scripts/setup_tests.py | 10 ++++ .evergreen/scripts/utils.py | 2 +- CONTRIBUTING.md | 13 ++++++ pyproject.toml | 2 +- test/atlas/test_connection.py | 2 +- 12 files changed, 91 insertions(+), 97 deletions(-) delete mode 100755 .evergreen/scripts/run-atlas-tests.sh delete mode 100755 .evergreen/scripts/run-direct-tests.sh delete mode 100755 .evergreen/scripts/run-enterprise-auth-tests.sh diff --git a/.evergreen/config.yml b/.evergreen/config.yml index f563a1ace..4562f1d2b 100644 --- a/.evergreen/config.yml +++ b/.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" diff --git a/.evergreen/generated_configs/tasks.yml b/.evergreen/generated_configs/tasks.yml index de18b4df0..0b0f09329 100644 --- a/.evergreen/generated_configs/tasks.yml +++ b/.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: diff --git a/.evergreen/generated_configs/variants.yml b/.evergreen/generated_configs/variants.yml index 938f98e43..4c54abf4b 100644 --- a/.evergreen/generated_configs/variants.yml +++ b/.evergreen/generated_configs/variants.yml @@ -49,7 +49,7 @@ 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 @@ -57,7 +57,7 @@ buildvariants: 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 diff --git a/.evergreen/scripts/generate_config.py b/.evergreen/scripts/generate_config.py index 84cd1075b..d91e0e6de 100644 --- a/.evergreen/scripts/generate_config.py +++ b/.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 = [ diff --git a/.evergreen/scripts/run-atlas-tests.sh b/.evergreen/scripts/run-atlas-tests.sh deleted file mode 100755 index 99968063b..000000000 --- a/.evergreen/scripts/run-atlas-tests.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/bin/bash - -# Disable xtrace for security reasons (just in case it was accidentally set). -set +x -set -o errexit -bash "${DRIVERS_TOOLS}"/.evergreen/auth_aws/setup_secrets.sh drivers/atlas_connect -bash "${PROJECT_DIRECTORY}"/.evergreen/just.sh setup-tests atlas -bash "${PROJECT_DIRECTORY}"/.evergreen/just.sh run-tests diff --git a/.evergreen/scripts/run-direct-tests.sh b/.evergreen/scripts/run-direct-tests.sh deleted file mode 100755 index a00235311..000000000 --- a/.evergreen/scripts/run-direct-tests.sh +++ /dev/null @@ -1,10 +0,0 @@ -#!/bin/bash -set -x -. .evergreen/utils.sh - -. .evergreen/scripts/env.sh -createvirtualenv "$PYTHON_BINARY" .venv - -export PYMONGO_C_EXT_MUST_BUILD=1 -pip install -e ".[test]" -pytest -v diff --git a/.evergreen/scripts/run-enterprise-auth-tests.sh b/.evergreen/scripts/run-enterprise-auth-tests.sh deleted file mode 100755 index 65aafde2d..000000000 --- a/.evergreen/scripts/run-enterprise-auth-tests.sh +++ /dev/null @@ -1,9 +0,0 @@ -#!/bin/bash -set -eu - -# Disable xtrace for security reasons (just in case it was accidentally set). -set +x -# Use the default python to bootstrap secrets. -bash "${DRIVERS_TOOLS}"/.evergreen/secrets_handling/setup-secrets.sh drivers/enterprise_auth -bash "${PROJECT_DIRECTORY}"/.evergreen/just.sh setup-tests enterprise_auth -bash "${PROJECT_DIRECTORY}"/.evergreen/just.sh run-tests diff --git a/.evergreen/scripts/setup_tests.py b/.evergreen/scripts/setup_tests.py index 3ba6d6175..868ac419b 100644 --- a/.evergreen/scripts/setup_tests.py +++ b/.evergreen/scripts/setup_tests.py @@ -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. diff --git a/.evergreen/scripts/utils.py b/.evergreen/scripts/utils.py index 039eec243..08d376461 100644 --- a/.evergreen/scripts/utils.py +++ b/.evergreen/scripts/utils.py @@ -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", diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7e70c025e..8844565d3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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`. diff --git a/pyproject.toml b/pyproject.toml index ca76cfa2c..993b3e5ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -125,7 +125,7 @@ markers = [ "auth_oidc: tests that rely on oidc auth", "auth: tests that rely on authentication", "ocsp: tests that rely on ocsp", - "atlas: tests that rely on atlas", + "atlas_connect: tests that rely on an atlas connection", "data_lake: tests that rely on atlas data lake", "perf: benchmark tests", "index_management: index management tests", diff --git a/test/atlas/test_connection.py b/test/atlas/test_connection.py index 4dcbba6d1..3d34ff326 100644 --- a/test/atlas/test_connection.py +++ b/test/atlas/test_connection.py @@ -28,7 +28,7 @@ sys.path[0:0] = [""] import pymongo from pymongo.ssl_support import HAS_SNI -pytestmark = pytest.mark.atlas +pytestmark = pytest.mark.atlas_connect URIS = { From e6e8650cc95314e742a30244b5db0239f5259f68 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Thu, 13 Mar 2025 15:08:41 -0400 Subject: [PATCH 2/2] PYTHON-5144 - Add async performance benchmarks (#2188) --- .evergreen/config.yml | 65 +++- .evergreen/run-perf-tests.sh | 2 +- .evergreen/scripts/run-perf-tests.sh | 2 +- .evergreen/scripts/setup_tests.py | 5 +- test/performance/async_perf_test.py | 466 +++++++++++++++++++++++++++ 5 files changed, 523 insertions(+), 17 deletions(-) create mode 100644 test/performance/async_perf_test.py diff --git a/.evergreen/config.yml b/.evergreen/config.yml index 4562f1d2b..24f08d67a 100644 --- a/.evergreen/config.yml +++ b/.evergreen/config.yml @@ -321,6 +321,7 @@ functions: params: working_dir: "src" binary: bash + include_expansions_in_env: [SUB_TEST_NAME] args: - .evergreen/scripts/run-with-env.sh - .evergreen/scripts/run-perf-tests.sh @@ -512,6 +513,8 @@ tasks: vars: VERSION: "v6.0-perf" - func: "run perf tests" + vars: + SUB_TEST_NAME: "sync" - func: "attach benchmark test results" - func: "send dashboard data" @@ -523,6 +526,8 @@ tasks: VERSION: "v6.0-perf" SSL: "ssl" - func: "run perf tests" + vars: + SUB_TEST_NAME: "sync" - func: "attach benchmark test results" - func: "send dashboard data" @@ -533,9 +538,52 @@ tasks: vars: VERSION: "8.0" - func: "run perf tests" + vars: + SUB_TEST_NAME: "sync" - func: "attach benchmark test results" - func: "send dashboard data" + - name: "perf-6.0-standalone-async" + tags: [ "perf" ] + commands: + - func: "run server" + vars: + VERSION: "v6.0-perf" + TOPOLOGY: "server" + - func: "run perf tests" + vars: + SUB_TEST_NAME: "async" + - func: "attach benchmark test results" + - func: "send dashboard data" + + - name: "perf-6.0-standalone-ssl-async" + tags: [ "perf" ] + commands: + - func: "run server" + vars: + VERSION: "v6.0-perf" + TOPOLOGY: "server" + SSL: "ssl" + - func: "run perf tests" + vars: + SUB_TEST_NAME: "async" + - func: "attach benchmark test results" + - func: "send dashboard data" + + - name: "perf-8.0-standalone-async" + tags: [ "perf" ] + commands: + - func: "run server" + vars: + VERSION: "8.0" + TOPOLOGY: "server" + - func: "run perf tests" + vars: + SUB_TEST_NAME: "async" + - func: "attach benchmark test results" + - func: "send dashboard data" + + - name: "check-import-time" tags: ["pr"] commands: @@ -616,17 +664,6 @@ buildvariants: - name: "perf-6.0-standalone" - name: "perf-6.0-standalone-ssl" - name: "perf-8.0-standalone" - - # Platform notes - # i386 builds of OpenSSL or Cyrus SASL are not available - # Debian 8.1 only supports MongoDB 3.4+ - # SUSE12 s390x is only supported by MongoDB 3.4+ - # No enterprise build for Archlinux, SSL not available - # RHEL 7.6 and RHEL 8.4 only supports 3.6+. - # RHEL 7 only supports 2.6+ - # RHEL 7.1 ppc64le is only supported by MongoDB 3.2+ - # RHEL 7.2 s390x is only supported by MongoDB 3.4+ - # Solaris MongoDB SSL builds are not available - # Darwin MongoDB SSL builds are not available for 2.6 - # SUSE12 x86_64 is only supported by MongoDB 3.2+ - # vim: set et sw=2 ts=2 : + - name: "perf-6.0-standalone-async" + - name: "perf-6.0-standalone-ssl-async" + - name: "perf-8.0-standalone-async" diff --git a/.evergreen/run-perf-tests.sh b/.evergreen/run-perf-tests.sh index 5e423caa2..cf88b9371 100755 --- a/.evergreen/run-perf-tests.sh +++ b/.evergreen/run-perf-tests.sh @@ -15,5 +15,5 @@ export OUTPUT_FILE="${PROJECT_DIRECTORY}/results.json" export PYTHON_BINARY=/opt/mongodbtoolchain/v4/bin/python3 -bash ./.evergreen/just.sh setup-tests perf +bash ./.evergreen/just.sh setup-tests perf "${SUB_TEST_NAME}" bash ./.evergreen/just.sh run-tests diff --git a/.evergreen/scripts/run-perf-tests.sh b/.evergreen/scripts/run-perf-tests.sh index 69a369fee..e1c1311d6 100755 --- a/.evergreen/scripts/run-perf-tests.sh +++ b/.evergreen/scripts/run-perf-tests.sh @@ -1,4 +1,4 @@ #!/bin/bash PROJECT_DIRECTORY=${PROJECT_DIRECTORY} -bash "${PROJECT_DIRECTORY}"/.evergreen/run-perf-tests.sh +SUB_TEST_NAME=${SUB_TEST_NAME} bash "${PROJECT_DIRECTORY}"/.evergreen/run-perf-tests.sh diff --git a/.evergreen/scripts/setup_tests.py b/.evergreen/scripts/setup_tests.py index 868ac419b..53e3a568b 100644 --- a/.evergreen/scripts/setup_tests.py +++ b/.evergreen/scripts/setup_tests.py @@ -359,7 +359,10 @@ def handle_test_env() -> None: if test_name == "perf": # PYTHON-4769 Run perf_test.py directly otherwise pytest's test collection negatively # affects the benchmark results. - TEST_ARGS = f"test/performance/perf_test.py {TEST_ARGS}" + if sub_test_name == "sync": + TEST_ARGS = f"test/performance/perf_test.py {TEST_ARGS}" + else: + TEST_ARGS = f"test/performance/async_perf_test.py {TEST_ARGS}" # Add coverage if requested. # Only cover CPython. PyPy reports suspiciously low coverage. diff --git a/test/performance/async_perf_test.py b/test/performance/async_perf_test.py new file mode 100644 index 000000000..2ceee45bf --- /dev/null +++ b/test/performance/async_perf_test.py @@ -0,0 +1,466 @@ +# Copyright 2015-present MongoDB, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Asynchronous Tests for the MongoDB Driver Performance Benchmarking Spec. + +See https://github.com/mongodb/specifications/blob/master/source/benchmarking/benchmarking.md + + +To set up the benchmarks locally:: + + python -m pip install simplejson + git clone --depth 1 https://github.com/mongodb/specifications.git + pushd specifications/source/benchmarking/data + tar xf extended_bson.tgz + tar xf parallel.tgz + tar xf single_and_multi_document.tgz + popd + export TEST_PATH="specifications/source/benchmarking/data" + export OUTPUT_FILE="results.json" + +Then to run all benchmarks quickly:: + + FASTBENCH=1 python test/performance/async_perf_test.py -v + +To run individual benchmarks quickly:: + + FASTBENCH=1 python test/performance/async_perf_test.py -v TestRunCommand TestFindManyAndEmptyCursor +""" +from __future__ import annotations + +import asyncio +import os +import sys +import tempfile +import time +import warnings +from typing import Any, List, Optional, Union + +import pytest + +try: + import simplejson as json +except ImportError: + import json # type: ignore[no-redef] + +sys.path[0:0] = [""] + +from test.asynchronous import AsyncPyMongoTestCase, async_client_context, unittest + +from bson import encode +from gridfs import AsyncGridFSBucket +from pymongo import ( + DeleteOne, + InsertOne, + ReplaceOne, +) + +pytestmark = pytest.mark.perf + +# Spec says to use at least 1 minute cumulative execution time and up to 100 iterations or 5 minutes but that +# makes the benchmarks too slow. Instead, we use at least 30 seconds and at most 60 seconds. +NUM_ITERATIONS = 100 +MIN_ITERATION_TIME = 30 +MAX_ITERATION_TIME = 120 +NUM_DOCS = 10000 +# When debugging or prototyping it's often useful to run the benchmarks locally, set FASTBENCH=1 to run quickly. +if bool(os.getenv("FASTBENCH")): + NUM_ITERATIONS = 2 + MIN_ITERATION_TIME = 1 + MAX_ITERATION_TIME = 30 + NUM_DOCS = 1000 + +TEST_PATH = os.environ.get( + "TEST_PATH", os.path.join(os.path.dirname(os.path.realpath(__file__)), os.path.join("data")) +) + +OUTPUT_FILE = os.environ.get("OUTPUT_FILE") + +result_data: List = [] + + +def tearDownModule(): + output = json.dumps(result_data, indent=4) + if OUTPUT_FILE: + with open(OUTPUT_FILE, "w") as opf: + opf.write(output) + else: + print(output) + + +class Timer: + def __enter__(self): + self.start = time.monotonic() + return self + + def __exit__(self, *args): + self.end = time.monotonic() + self.interval = self.end - self.start + + +async def concurrent(n_tasks, func): + tasks = [func() for _ in range(n_tasks)] + await asyncio.gather(*tasks) + + +class PerformanceTest: + dataset: str + data_size: int + fail: Any + n_tasks: int = 1 + did_init: bool = False + + async def asyncSetUp(self): + await async_client_context.init() + self.setup_time = time.monotonic() + + async def asyncTearDown(self): + duration = time.monotonic() - self.setup_time + # Remove "Test" so that TestFlatEncoding is reported as "FlatEncoding". + name = self.__class__.__name__[4:] + median = self.percentile(50) + megabytes_per_sec = (self.data_size * self.n_tasks) / median / 1000000 + print( + f"Completed {self.__class__.__name__} {megabytes_per_sec:.3f} MB/s, MEDIAN={self.percentile(50):.3f}s, " + f"total time={duration:.3f}s, iterations={len(self.results)}" + ) + result_data.append( + { + "info": { + "test_name": name, + "args": { + "tasks": self.n_tasks, + }, + }, + "metrics": [ + {"name": "megabytes_per_sec", "type": "MEDIAN", "value": megabytes_per_sec}, + ], + } + ) + + async def before(self): + pass + + async def do_task(self): + raise NotImplementedError + + async def after(self): + pass + + def percentile(self, percentile): + if hasattr(self, "results"): + sorted_results = sorted(self.results) + percentile_index = int(len(sorted_results) * percentile / 100) - 1 + return sorted_results[percentile_index] + else: + self.fail("Test execution failed") + return None + + async def runTest(self): + results = [] + start = time.monotonic() + i = 0 + while True: + i += 1 + await self.before() + with Timer() as timer: + if self.n_tasks == 1: + await self.do_task() + else: + await concurrent(self.n_tasks, self.do_task) + await self.after() + results.append(timer.interval) + duration = time.monotonic() - start + if duration > MIN_ITERATION_TIME and i >= NUM_ITERATIONS: + break + if i >= NUM_ITERATIONS: + break + if duration > MAX_ITERATION_TIME: + with warnings.catch_warnings(): + warnings.simplefilter("default") + warnings.warn( + f"{self.__class__.__name__} timed out after {MAX_ITERATION_TIME}s, completed {i}/{NUM_ITERATIONS} iterations." + ) + + break + + self.results = results + + +# SINGLE-DOC BENCHMARKS +class TestRunCommand(PerformanceTest, AsyncPyMongoTestCase): + data_size = len(encode({"hello": True})) * NUM_DOCS + + async def asyncSetUp(self): + await super().asyncSetUp() + self.client = async_client_context.client + await self.client.drop_database("perftest") + + async def do_task(self): + command = self.client.perftest.command + for _ in range(NUM_DOCS): + await command("hello", True) + + +class TestRunCommand8Tasks(TestRunCommand): + n_tasks = 8 + + +class TestRunCommand80Tasks(TestRunCommand): + n_tasks = 80 + + +class TestRunCommandUnlimitedTasks(TestRunCommand): + async def do_task(self): + command = self.client.perftest.command + await asyncio.gather(*[command("hello", True) for _ in range(NUM_DOCS)]) + + +class TestDocument(PerformanceTest): + async def asyncSetUp(self): + await super().asyncSetUp() + # Location of test data. + with open( # noqa: ASYNC101 + os.path.join(TEST_PATH, os.path.join("single_and_multi_document", self.dataset)) + ) as data: + self.document = json.loads(data.read()) + + self.client = async_client_context.client + await self.client.drop_database("perftest") + + async def asyncTearDown(self): + await super().asyncTearDown() + await self.client.drop_database("perftest") + + async def before(self): + self.corpus = await self.client.perftest.create_collection("corpus") + + async def after(self): + await self.client.perftest.drop_collection("corpus") + + +class FindTest(TestDocument): + dataset = "tweet.json" + + async def asyncSetUp(self): + await super().asyncSetUp() + self.data_size = len(encode(self.document)) * NUM_DOCS + documents = [self.document.copy() for _ in range(NUM_DOCS)] + self.corpus = self.client.perftest.corpus + result = await self.corpus.insert_many(documents) + self.inserted_ids = result.inserted_ids + + async def before(self): + pass + + async def after(self): + pass + + +class TestFindOneByID(FindTest, AsyncPyMongoTestCase): + async def do_task(self): + find_one = self.corpus.find_one + for _id in self.inserted_ids: + await find_one({"_id": _id}) + + +class TestFindOneByID8Tasks(TestFindOneByID): + n_tasks = 8 + + +class TestFindOneByID80Tasks(TestFindOneByID): + n_tasks = 80 + + +class TestFindOneByIDUnlimitedTasks(TestFindOneByID): + async def do_task(self): + find_one = self.corpus.find_one + await asyncio.gather(*[find_one({"_id": _id}) for _id in self.inserted_ids]) + + +class SmallDocInsertTest(TestDocument): + dataset = "small_doc.json" + + async def asyncSetUp(self): + await super().asyncSetUp() + self.data_size = len(encode(self.document)) * NUM_DOCS + self.documents = [self.document.copy() for _ in range(NUM_DOCS)] + + +class SmallDocMixedTest(TestDocument): + dataset = "small_doc.json" + + async def asyncSetUp(self): + await super().asyncSetUp() + self.data_size = len(encode(self.document)) * NUM_DOCS * 2 + self.documents = [self.document.copy() for _ in range(NUM_DOCS)] + + +class TestSmallDocInsertOne(SmallDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + insert_one = self.corpus.insert_one + for doc in self.documents: + await insert_one(doc) + + +class TestSmallDocInsertOneUnlimitedTasks(SmallDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + insert_one = self.corpus.insert_one + await asyncio.gather(*[insert_one(doc) for doc in self.documents]) + + +class LargeDocInsertTest(TestDocument): + dataset = "large_doc.json" + + async def asyncSetUp(self): + await super().asyncSetUp() + n_docs = 10 + self.data_size = len(encode(self.document)) * n_docs + self.documents = [self.document.copy() for _ in range(n_docs)] + + +class TestLargeDocInsertOne(LargeDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + insert_one = self.corpus.insert_one + for doc in self.documents: + await insert_one(doc) + + +class TestLargeDocInsertOneUnlimitedTasks(LargeDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + insert_one = self.corpus.insert_one + await asyncio.gather(*[insert_one(doc) for doc in self.documents]) + + +# MULTI-DOC BENCHMARKS +class TestFindManyAndEmptyCursor(FindTest, AsyncPyMongoTestCase): + async def do_task(self): + await self.corpus.find().to_list() + + +class TestFindManyAndEmptyCursor8Tasks(TestFindManyAndEmptyCursor): + n_tasks = 8 + + +class TestFindManyAndEmptyCursor80Tasks(TestFindManyAndEmptyCursor): + n_tasks = 80 + + +class TestSmallDocBulkInsert(SmallDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + await self.corpus.insert_many(self.documents, ordered=True) + + +class TestSmallDocClientBulkInsert(SmallDocInsertTest, AsyncPyMongoTestCase): + @async_client_context.require_version_min(8, 0, 0, -24) + async def asyncSetUp(self): + await super().asyncSetUp() + self.models = [] + for doc in self.documents: + self.models.append(InsertOne(namespace="perftest.corpus", document=doc)) + + @async_client_context.require_version_min(8, 0, 0, -24) + async def do_task(self): + await self.client.bulk_write(self.models, ordered=True) + + +class TestSmallDocBulkMixedOps(SmallDocMixedTest, AsyncPyMongoTestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.models: list[Union[InsertOne, ReplaceOne, DeleteOne]] = [] + for doc in self.documents: + self.models.append(InsertOne(document=doc)) + self.models.append(ReplaceOne(filter={}, replacement=doc.copy(), upsert=True)) + self.models.append(DeleteOne(filter={})) + + async def do_task(self): + await self.corpus.bulk_write(self.models, ordered=True) + + +class TestSmallDocClientBulkMixedOps(SmallDocMixedTest, AsyncPyMongoTestCase): + @async_client_context.require_version_min(8, 0, 0, -24) + async def asyncSetUp(self): + await super().asyncSetUp() + self.models: list[Union[InsertOne, ReplaceOne, DeleteOne]] = [] + for doc in self.documents: + self.models.append(InsertOne(namespace="perftest.corpus", document=doc)) + self.models.append( + ReplaceOne( + namespace="perftest.corpus", filter={}, replacement=doc.copy(), upsert=True + ) + ) + self.models.append(DeleteOne(namespace="perftest.corpus", filter={})) + + @async_client_context.require_version_min(8, 0, 0, -24) + async def do_task(self): + await self.client.bulk_write(self.models, ordered=True) + + +class TestLargeDocBulkInsert(LargeDocInsertTest, AsyncPyMongoTestCase): + async def do_task(self): + await self.corpus.insert_many(self.documents, ordered=True) + + +class TestLargeDocClientBulkInsert(LargeDocInsertTest, AsyncPyMongoTestCase): + @async_client_context.require_version_min(8, 0, 0, -24) + async def asyncSetUp(self): + await super().asyncSetUp() + self.models = [] + for doc in self.documents: + self.models.append(InsertOne(namespace="perftest.corpus", document=doc)) + + @async_client_context.require_version_min(8, 0, 0, -24) + async def do_task(self): + await self.client.bulk_write(self.models, ordered=True) + + +class GridFsTest(PerformanceTest): + async def asyncSetUp(self): + await super().asyncSetUp() + self.client = async_client_context.client + await self.client.drop_database("perftest") + + gridfs_path = os.path.join( + TEST_PATH, os.path.join("single_and_multi_document", "gridfs_large.bin") + ) + with open(gridfs_path, "rb") as data: # noqa: ASYNC101 + self.document = data.read() + self.data_size = len(self.document) + self.bucket = AsyncGridFSBucket(self.client.perftest) + + async def asyncTearDown(self): + await super().asyncTearDown() + await self.client.drop_database("perftest") + + +class TestGridFsUpload(GridFsTest, AsyncPyMongoTestCase): + async def before(self): + # Create the bucket. + await self.bucket.upload_from_stream("init", b"x") + + async def do_task(self): + await self.bucket.upload_from_stream("gridfstest", self.document) + + +class TestGridFsDownload(GridFsTest, AsyncPyMongoTestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.uploaded_id = await self.bucket.upload_from_stream("gridfstest", self.document) + + async def do_task(self): + await (await self.bucket.open_download_stream(self.uploaded_id)).read() + + +if __name__ == "__main__": + unittest.main()