diff --git a/test/performance/async_perf_test.py b/test/performance/async_perf_test.py index 6eb31ea4fe..961a2ddac2 100644 --- a/test/performance/async_perf_test.py +++ b/test/performance/async_perf_test.py @@ -56,11 +56,12 @@ sys.path[0:0] = [""] -from test.asynchronous import AsyncPyMongoTestCase, async_client_context, unittest +from test.asynchronous import AsyncPyMongoTestCase, async_client_context, get_loop, unittest from bson import encode from gridfs import AsyncGridFSBucket from pymongo import ( + AsyncMongoClient, DeleteOne, InsertOne, ReplaceOne, @@ -98,6 +99,9 @@ def tearDownModule(): else: print(output) + if getattr(async_client_context, "client", None): + get_loop().run_until_complete(async_client_context.client.close()) + class Timer: def __enter__(self): @@ -297,6 +301,56 @@ async def do_task(self): await asyncio.gather(*[find_one({"_id": _id}) for _id in self.inserted_ids]) +class SmallReadTest(PerformanceTest): + dataset = "small_doc.json" + n_tasks = 1 + + async def asyncSetUp(self): + await super().asyncSetUp() + 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.data_size = len(encode(self.document)) * NUM_DOCS + + client_options = dict(async_client_context.client_options) + client_options["minPoolSize"] = self.n_tasks + self.client = AsyncMongoClient(**client_options) + + await self.client.drop_database("perftest") + self.corpus = self.client.perftest.corpus + result = await self.corpus.insert_many([self.document.copy() for _ in range(NUM_DOCS)]) + self.inserted_ids = result.inserted_ids + + await self.client.admin.command("ping") + await self.corpus.find_one({"_id": self.inserted_ids[0]}) + + async def asyncTearDown(self): + try: + await super().asyncTearDown() + await self.client.drop_database("perftest") + finally: + await self.client.close() + + async def before(self): + pass + + async def after(self): + pass + + +class TestSmallReadFindOneByID(SmallReadTest, AsyncPyMongoTestCase): + async def do_task(self): + find_one = self.corpus.find_one + for _id in self.inserted_ids: + await find_one({"_id": _id}) + + +class TestSmallReadFindOneByID8Tasks(TestSmallReadFindOneByID): + n_tasks = 8 + + class SmallDocInsertTest(TestDocument): dataset = "small_doc.json" diff --git a/test/performance/perf_test.py b/test/performance/perf_test.py index 5688d28d2d..9d4c2400d5 100644 --- a/test/performance/perf_test.py +++ b/test/performance/perf_test.py @@ -100,6 +100,9 @@ def tearDownModule(): else: print(output) + if getattr(client_context, "client", None): + client_context.client.close() + class Timer: def __enter__(self): @@ -394,6 +397,74 @@ class TestFindOneByID8Threads(TestFindOneByID): n_threads = 8 +class SmallReadTest(PerformanceTest): + dataset = "small_doc.json" + n_threads = 1 + + def setUp(self): + super().setUp() + with open( + os.path.join(TEST_PATH, os.path.join("single_and_multi_document", self.dataset)) + ) as data: + self.document = json.loads(data.read()) + + self.data_size = len(encode(self.document)) * NUM_DOCS + + client_options = dict(client_context.client_options) + client_options["minPoolSize"] = self.n_threads + self.client = MongoClient(**client_options) + + self.client.drop_database("perftest") + self.corpus = self.client.perftest.corpus + result = self.corpus.insert_many([self.document.copy() for _ in range(NUM_DOCS)]) + self.inserted_ids = result.inserted_ids + + self.client.admin.command("ping") + self.corpus.find_one({"_id": self.inserted_ids[0]}) + + def tearDown(self): + try: + super().tearDown() + self.client.drop_database("perftest") + finally: + self.client.close() + + def before(self): + pass + + def after(self): + pass + + +class TestSmallReadFindOneByID(SmallReadTest, unittest.TestCase): + def do_task(self): + find_one = self.corpus.find_one + for _id in self.inserted_ids: + find_one({"_id": _id}) + + +class TestSmallReadFindOneByID8Threads(TestSmallReadFindOneByID): + n_threads = 8 + + +class TestSmallReadFindOneByID8ThreadsPerClient(SmallReadTest, unittest.TestCase): + n_threads = 8 + + def do_task(self): + client_options = dict(client_context.client_options) + client = MongoClient(**client_options) + try: + corpus = client.perftest.corpus + client.admin.command("ping") + corpus.find_one({"_id": self.inserted_ids[0]}) + + find_one = corpus.find_one + for _id in self.inserted_ids: + find_one({"_id": _id}) + finally: + client.close() + + class SmallDocInsertTest(TestDocument): dataset = "small_doc.json"