PYTHON-5113 - Refactor test utils for async (#2149)

Author: NoahStapp

Diff for: test/asynchronous/test_async_cancellation.py

@@ -17,7 +17,8 @@
 
 import asyncio
 import sys
-from test.utils import async_get_pool, delay, one
+from test.asynchronous.utils import async_get_pool
+from test.utils_shared import delay, one
 
 sys.path[0:0] = [""]
 

Diff for: test/asynchronous/test_auth.py

@@ -30,7 +30,7 @@
     async_client_context,
     unittest,
 )
-from test.utils import AllowListEventListener, delay, ignore_deprecations
+from test.utils_shared import AllowListEventListener, delay, ignore_deprecations
 
 import pytest
 

Diff for: test/asynchronous/test_bulk.py

@@ -24,7 +24,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, remove_all_users, unittest
-from test.utils import async_wait_until
+from test.utils_shared import async_wait_until
 
 from bson.binary import Binary, UuidRepresentation
 from bson.codec_options import CodecOptions

Diff for: test/asynchronous/test_change_stream.py

@@ -36,7 +36,7 @@
     unittest,
 )
 from test.asynchronous.unified_format import generate_test_classes
-from test.utils import (
+from test.utils_shared import (
     AllowListEventListener,
     EventListener,
     OvertCommandListener,

Diff for: test/asynchronous/test_client.py

@@ -60,14 +60,16 @@
     unittest,
 )
 from test.asynchronous.pymongo_mocks import AsyncMockClient
+from test.asynchronous.utils import (
+    async_get_pool,
+    async_wait_until,
+    asyncAssertRaisesExactly,
+)
 from test.test_binary import BinaryData
-from test.utils import (
+from test.utils_shared import (
     NTHREADS,
     CMAPListener,
     FunctionCallRecorder,
-    async_get_pool,
-    async_wait_until,
-    asyncAssertRaisesExactly,
     delay,
     gevent_monkey_patched,
     is_greenthread_patched,

Diff for: test/asynchronous/test_client_bulk_write.py

@@ -25,7 +25,7 @@
     async_client_context,
     unittest,
 )
-from test.utils import (
+from test.utils_shared import (
     OvertCommandListener,
 )
 from unittest.mock import patch

Diff for: test/asynchronous/test_collation.py

@@ -18,7 +18,7 @@
 import functools
 import warnings
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import EventListener, OvertCommandListener
+from test.utils_shared import EventListener, OvertCommandListener
 from typing import Any
 
 from pymongo.asynchronous.helpers import anext

Diff for: test/asynchronous/test_collection.py

@@ -21,6 +21,7 @@
 import sys
 from codecs import utf_8_decode
 from collections import defaultdict
+from test.asynchronous.utils import async_get_pool, async_is_mongos
 from typing import Any, Iterable, no_type_check
 
 from pymongo.asynchronous.database import AsyncDatabase
@@ -33,12 +34,10 @@
     AsyncUnitTest,
     async_client_context,
 )
-from test.utils import (
+from test.utils_shared import (
     IMPOSSIBLE_WRITE_CONCERN,
     EventListener,
     OvertCommandListener,
-    async_get_pool,
-    async_is_mongos,
     async_wait_until,
 )
 

Diff for: test/asynchronous/test_comment.py

@@ -22,7 +22,7 @@
 sys.path[0:0] = [""]
 from asyncio import iscoroutinefunction
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import OvertCommandListener
+from test.utils_shared import OvertCommandListener
 
 from bson.dbref import DBRef
 from pymongo.asynchronous.command_cursor import AsyncCommandCursor

Diff for: test/asynchronous/test_concurrency.py

@@ -18,7 +18,7 @@
 import asyncio
 import time
 from test.asynchronous import AsyncIntegrationTest, async_client_context
-from test.utils import delay
+from test.utils_shared import delay
 
 _IS_SYNC = False
 

Diff for: test/asynchronous/test_connection_monitoring.py

@@ -20,17 +20,15 @@
 import sys
 import time
 from pathlib import Path
+from test.asynchronous.utils import async_get_pool, async_get_pools
 
 sys.path[0:0] = [""]
 
-from test.asynchronous import AsyncIntegrationTest, client_knobs, unittest
+from test.asynchronous import AsyncIntegrationTest, async_client_context, client_knobs, unittest
 from test.asynchronous.pymongo_mocks import DummyMonitor
 from test.asynchronous.utils_spec_runner import AsyncSpecTestCreator, SpecRunnerTask
-from test.utils import (
+from test.utils_shared import (
     CMAPListener,
-    async_client_context,
-    async_get_pool,
-    async_get_pools,
     async_wait_until,
     camel_to_snake,
 )

Diff for: test/asynchronous/test_connections_survive_primary_stepdown_spec.py

@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import sys
+from test.asynchronous.utils import async_ensure_all_connected
 
 sys.path[0:0] = [""]
 
@@ -25,9 +26,8 @@
     unittest,
 )
 from test.asynchronous.helpers import async_repl_set_step_down
-from test.utils import (
+from test.utils_shared import (
     CMAPListener,
-    async_ensure_all_connected,
 )
 
 from bson import SON

Diff for: test/asynchronous/test_cursor.py

@@ -31,7 +31,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import (
+from test.utils_shared import (
     AllowListEventListener,
     EventListener,
     OvertCommandListener,

Diff for: test/asynchronous/test_data_lake.py

@@ -25,7 +25,7 @@
 
 from test.asynchronous import AsyncIntegrationTest, AsyncUnitTest, async_client_context, unittest
 from test.asynchronous.unified_format import generate_test_classes
-from test.utils import (
+from test.utils_shared import (
     OvertCommandListener,
 )
 

Diff for: test/asynchronous/test_database.py

@@ -26,7 +26,7 @@
 from test import unittest
 from test.asynchronous import AsyncIntegrationTest, async_client_context
 from test.test_custom_types import DECIMAL_CODECOPTS
-from test.utils import (
+from test.utils_shared import (
     IMPOSSIBLE_WRITE_CONCERN,
     OvertCommandListener,
     async_wait_until,

Diff for: test/asynchronous/test_discovery_and_monitoring.py

@@ -26,25 +26,32 @@
 
 sys.path[0:0] = [""]
 
-from test.asynchronous import AsyncIntegrationTest, AsyncPyMongoTestCase, AsyncUnitTest, unittest
+from test.asynchronous import (
+    AsyncIntegrationTest,
+    AsyncPyMongoTestCase,
+    AsyncUnitTest,
+    async_client_context,
+    unittest,
+)
 from test.asynchronous.pymongo_mocks import DummyMonitor
 from test.asynchronous.unified_format import generate_test_classes
-from test.utils import (
+from test.asynchronous.utils import (
+    async_get_pool,
+)
+from test.utils_shared import (
     CMAPListener,
     HeartbeatEventListener,
     HeartbeatEventsListListener,
     assertion_context,
     async_barrier_wait,
-    async_client_context,
     async_create_barrier,
-    async_get_pool,
     async_wait_until,
     server_name_to_type,
 )
 from unittest.mock import patch
 
 from bson import Timestamp, json_util
-from pymongo import AsyncMongoClient, common, monitoring
+from pymongo import common, monitoring
 from pymongo.asynchronous.settings import TopologySettings
 from pymongo.asynchronous.topology import Topology, _ErrorContext
 from pymongo.errors import (
@@ -291,7 +298,7 @@ async def test_ignore_stale_connection_errors(self):
         if not _IS_SYNC and sys.version_info < (3, 11):
             self.skipTest("Test requires asyncio.Barrier (added in Python 3.11)")
         N_TASKS = 5
-        barrier = async_create_barrier(N_TASKS, timeout=30)
+        barrier = async_create_barrier(N_TASKS)
         client = await self.async_rs_or_single_client(minPoolSize=N_TASKS)
 
         # Wait for initial discovery.

Diff for: test/asynchronous/test_dns.py

@@ -29,7 +29,7 @@
     async_client_context,
     unittest,
 )
-from test.utils import async_wait_until
+from test.utils_shared import async_wait_until
 
 from pymongo.common import validate_read_preference_tags
 from pymongo.errors import ConfigurationError

Diff for: test/asynchronous/test_encryption.py

@@ -64,7 +64,7 @@
     KMIP_CREDS,
     LOCAL_MASTER_KEY,
 )
-from test.utils import (
+from test.utils_shared import (
     AllowListEventListener,
     OvertCommandListener,
     TopologyEventListener,

Diff for: test/asynchronous/test_examples.py

@@ -26,7 +26,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import async_wait_until
+from test.utils_shared import async_wait_until
 
 import pymongo
 from pymongo.asynchronous.helpers import anext

Diff for: test/asynchronous/test_grid_file.py

@@ -33,7 +33,7 @@
 
 sys.path[0:0] = [""]
 
-from test.utils import OvertCommandListener
+from test.utils_shared import OvertCommandListener
 
 from bson.objectid import ObjectId
 from gridfs.asynchronous.grid_file import (

Diff for: test/asynchronous/test_gridfs.py

@@ -28,7 +28,8 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import async_joinall, one
+from test.asynchronous.utils import async_joinall
+from test.utils_shared import one
 
 import gridfs
 from bson.binary import Binary

Diff for: test/asynchronous/test_gridfs_bucket.py

@@ -29,7 +29,8 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.utils import async_joinall, joinall, one
+from test.asynchronous.utils import async_joinall
+from test.utils_shared import one
 
 import gridfs
 from bson.binary import Binary

Diff for: test/asynchronous/test_heartbeat_monitoring.py

@@ -16,11 +16,12 @@
 from __future__ import annotations
 
 import sys
+from test.asynchronous.utils import AsyncMockPool
 
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, client_knobs, unittest
-from test.utils import AsyncMockPool, HeartbeatEventListener, async_wait_until
+from test.utils_shared import HeartbeatEventListener, async_wait_until
 
 from pymongo.asynchronous.monitor import Monitor
 from pymongo.errors import ConnectionFailure

Diff for: test/asynchronous/test_index_management.py

@@ -29,7 +29,7 @@
 
 from test.asynchronous import AsyncIntegrationTest, AsyncPyMongoTestCase, unittest
 from test.asynchronous.unified_format import generate_test_classes
-from test.utils import AllowListEventListener, OvertCommandListener
+from test.utils_shared import AllowListEventListener, OvertCommandListener
 
 from pymongo.errors import OperationFailure
 from pymongo.operations import SearchIndexModel

Diff for: test/asynchronous/test_load_balancer.py

@@ -23,15 +23,15 @@
 import threading
 from asyncio import Event
 from test.asynchronous.helpers import ConcurrentRunner, ExceptionCatchingTask
+from test.asynchronous.utils import async_get_pool
 
 import pytest
 
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
 from test.asynchronous.unified_format import generate_test_classes
-from test.utils import (
-    async_get_pool,
+from test.utils_shared import (
     async_wait_until,
     create_async_event,
 )

Diff for: test/asynchronous/test_max_staleness.py

@@ -28,7 +28,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncPyMongoTestCase, async_client_context, unittest
-from test.utils_selection_tests import create_selection_tests
+from test.asynchronous.utils_selection_tests import create_selection_tests
 
 from pymongo.errors import ConfigurationError
 from pymongo.server_selectors import writable_server_selector

Diff for: test/asynchronous/test_mongos_load_balancing.py

@@ -26,7 +26,7 @@
 
 from test.asynchronous import AsyncMockClientTest, async_client_context, connected, unittest
 from test.asynchronous.pymongo_mocks import AsyncMockClient
-from test.utils import async_wait_until
+from test.utils_shared import async_wait_until
 
 from pymongo.errors import AutoReconnect, InvalidOperation
 from pymongo.server_selectors import writable_server_selector

Diff for: test/asynchronous/test_monitor.py

@@ -25,10 +25,10 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, connected, unittest
-from test.utils import (
-    ServerAndTopologyEventListener,
+from test.asynchronous.utils import (
     async_wait_until,
 )
+from test.utils_shared import ServerAndTopologyEventListener
 
 from pymongo.periodic_executor import _EXECUTORS
 

Diff for: test/asynchronous/test_monitoring.py

@@ -29,7 +29,7 @@
     sanitize_cmd,
     unittest,
 )
-from test.utils import (
+from test.utils_shared import (
     EventListener,
     OvertCommandListener,
     async_wait_until,

Diff for: test/asynchronous/test_pooling.py

@@ -21,6 +21,7 @@
 import socket
 import sys
 import time
+from test.asynchronous.utils import async_get_pool, async_joinall
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
 from bson.son import SON
@@ -33,7 +34,7 @@
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
 from test.asynchronous.helpers import ConcurrentRunner
-from test.utils import async_get_pool, async_joinall, delay
+from test.utils_shared import delay
 
 from pymongo.asynchronous.pool import Pool, PoolOptions
 from pymongo.socket_checker import SocketChecker

Diff for: test/asynchronous/test_read_concern.py

@@ -21,7 +21,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context
-from test.utils import OvertCommandListener
+from test.utils_shared import OvertCommandListener
 
 from bson.son import SON
 from pymongo.errors import OperationFailure