mongodb/mongo-python-driver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'master' of github.com:mongodb/mongo-python-driver

Browse Source
This commit is contained in:
Steven Silvester 2025-03-13 16:52:48 -05:00
commit 9bd25fb6c6
No known key found for this signature in database
GPG Key ID: B1BF5EC3A8B32F91
15 changed files with 614 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

100
.evergreen/config.yml
View File

@ -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:
@ -343,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
@ -422,13 +401,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 +460,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"
@ -547,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"
@ -558,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"
@ -568,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:
@ -651,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"

23
.evergreen/generated_configs/tasks.yml
View File

@ -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:

46
.evergreen/generated_configs/variants.yml
View File

@ -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

2
.evergreen/run-perf-tests.sh
View File

@ -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

28
.evergreen/scripts/generate_config.py
View File

@ -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 = [

8
.evergreen/scripts/run-atlas-tests.sh
View File

@ -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

10
.evergreen/scripts/run-direct-tests.sh
View File

@ -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

9
.evergreen/scripts/run-enterprise-auth-tests.sh
View File

@ -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

2
.evergreen/scripts/run-perf-tests.sh
View File

@ -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

15
.evergreen/scripts/setup_tests.py
View File

@ -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,10 +351,18 @@ 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.
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.

2
.evergreen/scripts/utils.py
View File

@ -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",

13
CONTRIBUTING.md
View File

@ -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`.

2
pyproject.toml
View File

@ -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",

2
test/atlas/test_connection.py
View File

@ -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 = {

466
test/performance/async_perf_test.py Normal file
View File

@ -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()