diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py index 1ad3e66d00..9834cc87c4 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 ConcurrentRunner from test.asynchronous.pymongo_mocks import AsyncMockClient from test.asynchronous.utils import ( async_get_pool, @@ -1980,27 +1980,17 @@ 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 + running = True - def stop(self): - self.running = False + 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) - 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) - - def run(self): - self._run() - - t = ResetPoolThread(pool) - t.start() + t = ConcurrentRunner(target=reset_pool) + await t.start() # Ensure that update_pool completes without error even when the pool # is reset concurrently. @@ -2011,8 +2001,8 @@ def run(self): if generation != pool.gen.get_overall(): break finally: - t.stop() - t.join() + running = False + await t.join() 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..e801836eb7 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 @@ -83,17 +84,19 @@ def succeeded(self, event: CommandSucceededEvent) -> None: 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: + create_task(async_client_context.client.admin.command(cmd)) # type: ignore[arg-type] def retryable_single_statement_ops(coll): diff --git a/test/test_client.py b/test/test_client.py index 35eed1f219..f416ceb859 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 ConcurrentRunner from test.pymongo_mocks import MockClient from test.test_binary import BinaryData from test.utils import ( @@ -1933,26 +1933,16 @@ 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 = ConcurrentRunner(target=reset_pool) t.start() # Ensure that update_pool completes without error even when the pool @@ -1964,7 +1954,7 @@ def run(self): if generation != pool.gen.get_overall(): break finally: - t.stop() + running = False t.join() client.admin.command("ping") diff --git a/test/test_retryable_writes.py b/test/test_retryable_writes.py index 71552b208e..b8933a488d 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 @@ -83,17 +84,19 @@ def succeeded(self, event: CommandSucceededEvent) -> None: 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: + create_task(client_context.client.admin.command(cmd)) # type: ignore[arg-type] def retryable_single_statement_ops(coll):