diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 1ad3e66d00..f6fcde4247 100644 --- a/test/asynchronous/test_client.py +++ b/test/asynchronous/test_client.py @@ -31,7 +31,6 @@ import struct import subprocess import sys -import threading import time import uuid from collections.abc import Iterable @@ -113,6 +112,7 @@ remove_all_users, unittest, ) +from test.asynchronous.helpers import ExceptionCatchingTask from test.asynchronous.pymongo_mocks import AsyncMockClient from test.asynchronous.utils import ( async_get_pool, @@ -1972,7 +1972,6 @@ def decompress_spy( if not negotiated: self.skipTest("server did not negotiate compression for any compressor") - @async_client_context.require_sync async def test_reset_during_update_pool(self): client = await self.async_rs_or_single_client(minPoolSize=10) await client.admin.command("ping") @@ -1980,39 +1979,30 @@ async def test_reset_during_update_pool(self): generation = pool.gen.get_overall() # Continuously reset the pool. - class ResetPoolThread(threading.Thread): - def __init__(self, pool): - super().__init__() - self.running = True - self.pool = pool - - def stop(self): - self.running = False + running = True - async def _run(self): - while self.running: - exc = AutoReconnect("mock pool error") - ctx = _ErrorContext(exc, 0, pool.gen.get_overall(), False, None) - await client._topology.handle_error(pool.address, ctx) - await asyncio.sleep(0.001) + async def reset_pool(): + while running: + exc = AutoReconnect("mock pool error") + ctx = _ErrorContext(exc, 0, pool.gen.get_overall(), False, None) + await client._topology.handle_error(pool.address, ctx) + await asyncio.sleep(0.001) - def run(self): - self._run() - - t = ResetPoolThread(pool) - t.start() + t = ExceptionCatchingTask(target=reset_pool) + await t.start() # Ensure that update_pool completes without error even when the pool # is reset concurrently. try: - while True: + while t.is_alive(): for _ in range(10): await client._topology.update_pool() if generation != pool.gen.get_overall(): break finally: - t.stop() - t.join() + running = False + await t.join() + self.assertIsNone(t.exc) await client.admin.command("ping") async def test_background_connections_do_not_hold_locks(self): diff --git a/test/asynchronous/test_retryable_writes.py b/test/asynchronous/test_retryable_writes.py index 90be59016a..22eb014cd9 100644 --- a/test/asynchronous/test_retryable_writes.py +++ b/test/asynchronous/test_retryable_writes.py @@ -23,6 +23,7 @@ import threading from unittest import mock +from pymongo._asyncio_task import create_task from pymongo.common import MAX_ADAPTIVE_RETRIES from test.asynchronous.utils import async_set_fail_point, flaky @@ -77,23 +78,34 @@ class InsertEventListener(EventListener): + fail_point_exc: BaseException | None = None + + def _store_fail_point_exc(self, task: asyncio.Task) -> None: + if not task.cancelled(): + self.fail_point_exc = task.exception() + def succeeded(self, event: CommandSucceededEvent) -> None: super().succeeded(event) if ( event.command_name == "insert" and event.reply.get("writeConcernError", {}).get("code", None) == 91 ): - async_client_context.client.admin.command( - { - "configureFailPoint": "failCommand", - "mode": {"times": 1}, - "data": { - "errorCode": 10107, - "errorLabels": ["RetryableWriteError", "NoWritesPerformed"], - "failCommands": ["insert"], - }, - } - ) + cmd = { + "configureFailPoint": "failCommand", + "mode": {"times": 1}, + "data": { + "errorCode": 10107, + "errorLabels": ["RetryableWriteError", "NoWritesPerformed"], + "failCommands": ["insert"], + }, + } + if _IS_SYNC: + async_client_context.client.admin.command(cmd) + else: + # succeeded() cannot await, so the fail point may be configured after + # the driver dispatches the retry it is meant to fail. + self._task = create_task(async_client_context.client.admin.command(cmd)) # type: ignore[arg-type] + self._task.add_done_callback(self._store_fail_point_exc) def retryable_single_statement_ops(coll): @@ -559,6 +571,7 @@ async def test_returns_original_error_code( with self.assertRaises(WriteConcernError) as exc: await client.test.test.insert_one({"_id": 1}) self.assertEqual(exc.exception.code, 91) + self.assertIsNone(cmd_listener.fail_point_exc) await client.admin.command( { "configureFailPoint": "failCommand", diff --git a/test/test_client.py b/test/test_client.py index 35eed1f219..bde8c9b325 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -31,7 +31,6 @@ import struct import subprocess import sys -import threading import time import uuid from collections.abc import Iterable @@ -112,6 +111,7 @@ remove_all_users, unittest, ) +from test.helpers import ExceptionCatchingTask from test.pymongo_mocks import MockClient from test.test_binary import BinaryData from test.utils import ( @@ -1925,7 +1925,6 @@ def decompress_spy( if not negotiated: self.skipTest("server did not negotiate compression for any compressor") - @client_context.require_sync def test_reset_during_update_pool(self): client = self.rs_or_single_client(minPoolSize=10) client.admin.command("ping") @@ -1933,39 +1932,30 @@ def test_reset_during_update_pool(self): generation = pool.gen.get_overall() # Continuously reset the pool. - class ResetPoolThread(threading.Thread): - def __init__(self, pool): - super().__init__() - self.running = True - self.pool = pool - - def stop(self): - self.running = False - - def _run(self): - while self.running: - exc = AutoReconnect("mock pool error") - ctx = _ErrorContext(exc, 0, pool.gen.get_overall(), False, None) - client._topology.handle_error(pool.address, ctx) - time.sleep(0.001) - - def run(self): - self._run() - - t = ResetPoolThread(pool) + running = True + + def reset_pool(): + while running: + exc = AutoReconnect("mock pool error") + ctx = _ErrorContext(exc, 0, pool.gen.get_overall(), False, None) + client._topology.handle_error(pool.address, ctx) + time.sleep(0.001) + + t = ExceptionCatchingTask(target=reset_pool) t.start() # Ensure that update_pool completes without error even when the pool # is reset concurrently. try: - while True: + while t.is_alive(): for _ in range(10): client._topology.update_pool() if generation != pool.gen.get_overall(): break finally: - t.stop() + running = False t.join() + self.assertIsNone(t.exc) client.admin.command("ping") def test_background_connections_do_not_hold_locks(self): diff --git a/test/test_retryable_writes.py b/test/test_retryable_writes.py index 71552b208e..0746e22b26 100644 --- a/test/test_retryable_writes.py +++ b/test/test_retryable_writes.py @@ -23,6 +23,7 @@ import threading from unittest import mock +from pymongo._asyncio_task import create_task from pymongo.common import MAX_ADAPTIVE_RETRIES from test.utils import flaky, set_fail_point @@ -77,23 +78,34 @@ class InsertEventListener(EventListener): + fail_point_exc: BaseException | None = None + + def _store_fail_point_exc(self, task: asyncio.Task) -> None: + if not task.cancelled(): + self.fail_point_exc = task.exception() + def succeeded(self, event: CommandSucceededEvent) -> None: super().succeeded(event) if ( event.command_name == "insert" and event.reply.get("writeConcernError", {}).get("code", None) == 91 ): - client_context.client.admin.command( - { - "configureFailPoint": "failCommand", - "mode": {"times": 1}, - "data": { - "errorCode": 10107, - "errorLabels": ["RetryableWriteError", "NoWritesPerformed"], - "failCommands": ["insert"], - }, - } - ) + cmd = { + "configureFailPoint": "failCommand", + "mode": {"times": 1}, + "data": { + "errorCode": 10107, + "errorLabels": ["RetryableWriteError", "NoWritesPerformed"], + "failCommands": ["insert"], + }, + } + if _IS_SYNC: + client_context.client.admin.command(cmd) + else: + # succeeded() cannot await, so the fail point may be configured after + # the driver dispatches the retry it is meant to fail. + self._task = create_task(client_context.client.admin.command(cmd)) # type: ignore[arg-type] + self._task.add_done_callback(self._store_fail_point_exc) def retryable_single_statement_ops(coll): @@ -555,6 +567,7 @@ def test_returns_original_error_code( with self.assertRaises(WriteConcernError) as exc: client.test.test.insert_one({"_id": 1}) self.assertEqual(exc.exception.code, 91) + self.assertIsNone(cmd_listener.fail_point_exc) client.admin.command( { "configureFailPoint": "failCommand",