Fix failures

Author: NoahStapp

test/asynchronous/test_async_cancellation.py

@@ -17,7 +17,8 @@
 
 import asyncio
 import sys
-from test.utils_shared 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] = [""]
 

test/asynchronous/test_client.py

@@ -61,7 +61,6 @@
 )
 from test.asynchronous.pymongo_mocks import AsyncMockClient
 from test.asynchronous.utils import (
-    AsyncFunctionCallRecorder,
     async_get_pool,
     async_wait_until,
     asyncAssertRaisesExactly,
@@ -70,6 +69,7 @@
 from test.utils_shared import (
     NTHREADS,
     CMAPListener,
+    FunctionCallRecorder,
     delay,
     gevent_monkey_patched,
     is_greenthread_patched,
@@ -513,7 +513,7 @@ async def test_connection_timeout_ms_propagates_to_DNS_resolver(self):
         # Patch the resolver.
         from pymongo.srv_resolver import _resolve
 
-        patched_resolver = AsyncFunctionCallRecorder(_resolve)
+        patched_resolver = FunctionCallRecorder(_resolve)
         pymongo.srv_resolver._resolve = patched_resolver
 
         def reset_resolver():

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
@@ -37,8 +38,6 @@
     IMPOSSIBLE_WRITE_CONCERN,
     EventListener,
     OvertCommandListener,
-    async_get_pool,
-    async_is_mongos,
     async_wait_until,
 )
 

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_shared import (
     CMAPListener,
-    async_client_context,
-    async_get_pool,
-    async_get_pools,
     async_wait_until,
     camel_to_snake,
 )

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] = [""]
 
@@ -27,7 +28,6 @@
 from test.asynchronous.helpers import async_repl_set_step_down
 from test.utils_shared import (
     CMAPListener,
-    async_ensure_all_connected,
 )
 
 from bson import SON

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_shared import async_joinall, one
+from test.asynchronous.utils import async_joinall
+from test.utils_shared import one
 
 import gridfs
 from bson.binary import Binary

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

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

test/asynchronous/test_load_balancer.py

@@ -23,6 +23,7 @@
 import threading
 from asyncio import Event
 from test.asynchronous.helpers import ConcurrentRunner, ExceptionCatchingTask
+from test.asynchronous.utils import async_get_pool
 
 import pytest
 
@@ -31,7 +32,6 @@
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
 from test.asynchronous.unified_format import generate_test_classes
 from test.utils_shared import (
-    async_get_pool,
     async_wait_until,
     create_async_event,
 )

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

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

test/asynchronous/test_retryable_reads.py

@@ -19,6 +19,7 @@
 import pprint
 import sys
 import threading
+from test.asynchronous.utils import async_set_fail_point
 
 from pymongo.errors import AutoReconnect
 
@@ -34,7 +35,6 @@
 from test.utils_shared import (
     CMAPListener,
     OvertCommandListener,
-    async_set_fail_point,
 )
 
 from pymongo.monitoring import (

test/asynchronous/test_retryable_writes.py

@@ -20,6 +20,7 @@
 import pprint
 import sys
 import threading
+from test.asynchronous.utils import async_set_fail_point
 
 sys.path[0:0] = [""]
 
@@ -35,7 +36,6 @@
     DeprecationFilter,
     EventListener,
     OvertCommandListener,
-    async_set_fail_point,
 )
 from test.version import Version
 

test/asynchronous/test_server_selection.py

@@ -31,7 +31,7 @@
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncIntegrationTest, async_client_context, unittest
-from test.asynchronous.utils import AsyncFunctionCallRecorder, async_wait_until
+from test.asynchronous.utils import async_wait_until
 from test.asynchronous.utils_selection_tests import (
     create_selection_tests,
     get_topology_settings_dict,
@@ -41,6 +41,7 @@
     make_server_description,
 )
 from test.utils_shared import (
+    FunctionCallRecorder,
     OvertCommandListener,
 )
 
@@ -122,7 +123,7 @@ async def test_invalid_server_selector(self):
 
     @async_client_context.require_replica_set
     async def test_selector_called(self):
-        selector = AsyncFunctionCallRecorder(lambda x: x)
+        selector = FunctionCallRecorder(lambda x: x)
 
         # Client setup.
         mongo_client = await self.async_rs_or_single_client(server_selector=selector)
@@ -175,7 +176,7 @@ async def test_latency_threshold_application(self):
 
     @async_client_context.require_replica_set
     async def test_server_selector_bypassed(self):
-        selector = AsyncFunctionCallRecorder(lambda x: x)
+        selector = FunctionCallRecorder(lambda x: x)
 
         scenario_def = {
             "topology_description": {

test/asynchronous/test_server_selection_in_window.py

@@ -26,7 +26,6 @@
 from test.utils_shared import (
     CMAPListener,
     OvertCommandListener,
-    async_get_pool,
     async_wait_until,
 )
 

test/asynchronous/test_srv_polling.py

@@ -18,12 +18,13 @@
 import asyncio
 import sys
 import time
+from test.utils_shared import FunctionCallRecorder
 from typing import Any
 
 sys.path[0:0] = [""]
 
 from test.asynchronous import AsyncPyMongoTestCase, client_knobs, unittest
-from test.asynchronous.utils import AsyncFunctionCallRecorder, async_wait_until
+from test.asynchronous.utils import async_wait_until
 
 import pymongo
 from pymongo import common
@@ -69,7 +70,7 @@ def mock_get_hosts_and_min_ttl(resolver, *args):
 
         patch_func: Any
         if self.count_resolver_calls:
-            patch_func = AsyncFunctionCallRecorder(mock_get_hosts_and_min_ttl)
+            patch_func = FunctionCallRecorder(mock_get_hosts_and_min_ttl)
         else:
             patch_func = mock_get_hosts_and_min_ttl
 

test/asynchronous/unified_format.py

@@ -35,6 +35,7 @@
     client_knobs,
     unittest,
 )
+from test.asynchronous.utils import async_get_pool
 from test.asynchronous.utils_spec_runner import SpecRunnerTask
 from test.unified_format_shared import (
     KMS_TLS_OPTS,
@@ -50,7 +51,6 @@
     with_metaclass,
 )
 from test.utils_shared import (
-    async_get_pool,
     async_wait_until,
     camel_to_snake,
     camel_to_snake_args,

test/asynchronous/utils.py

@@ -17,57 +17,19 @@
 
 import asyncio
 import contextlib
-import copy
-import functools
-import os
 import random
-import re
-import shutil
-import sys
-import threading
 import time
-import unittest
-import warnings
 from asyncio import iscoroutinefunction
-from collections import abc, defaultdict
-from functools import partial
-from test import client_context, db_pwd, db_user
-from test.asynchronous import async_client_context
-from typing import Any, List
-
-from bson import json_util
-from bson.objectid import ObjectId
+
 from bson.son import SON
-from pymongo import AsyncMongoClient, monitoring, operations, read_preferences
-from pymongo._asyncio_task import create_task
-from pymongo.cursor_shared import CursorType
-from pymongo.errors import ConfigurationError, OperationFailure
+from pymongo import AsyncMongoClient
+from pymongo.errors import ConfigurationError
 from pymongo.hello import HelloCompat
-from pymongo.helpers_shared import _SENSITIVE_COMMANDS
-from pymongo.lock import _async_create_lock, _create_lock
-from pymongo.monitoring import (
-    ConnectionCheckedInEvent,
-    ConnectionCheckedOutEvent,
-    ConnectionCheckOutFailedEvent,
-    ConnectionCheckOutStartedEvent,
-    ConnectionClosedEvent,
-    ConnectionCreatedEvent,
-    ConnectionReadyEvent,
-    PoolClearedEvent,
-    PoolClosedEvent,
-    PoolCreatedEvent,
-    PoolReadyEvent,
-)
+from pymongo.lock import _async_create_lock
 from pymongo.operations import _Op
-from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_selectors import any_server_selector, writable_server_selector
-from pymongo.server_type import SERVER_TYPE
-from pymongo.synchronous.collection import ReturnDocument
-from pymongo.synchronous.mongo_client import MongoClient
 from pymongo.synchronous.pool import _CancellationContext, _PoolGeneration
-from pymongo.uri_parser import parse_uri
-from pymongo.write_concern import WriteConcern
 
 _IS_SYNC = False
 
@@ -246,31 +208,3 @@ async def update_is_writable(self, is_writable):
 
     async def remove_stale_sockets(self, *args, **kwargs):
         pass
-
-
-class AsyncFunctionCallRecorder:
-    """Utility class to wrap a callable and record its invocations."""
-
-    def __init__(self, function):
-        self._function = function
-        self._call_list = []
-
-    async def __call__(self, *args, **kwargs):
-        self._call_list.append((args, kwargs))
-        if iscoroutinefunction(self._function):
-            return await self._function(*args, **kwargs)
-        else:
-            return self._function(*args, **kwargs)
-
-    def reset(self):
-        """Wipes the call list."""
-        self._call_list = []
-
-    def call_list(self):
-        """Returns a copy of the call list."""
-        return self._call_list[:]
-
-    @property
-    def call_count(self):
-        """Returns the number of times the function has been called."""
-        return len(self._call_list)

test/asynchronous/utils_selection_tests.py

@@ -19,17 +19,18 @@
 import os
 import sys
 from test.asynchronous import AsyncPyMongoTestCase
+from test.asynchronous.utils import AsyncMockPool
 
 sys.path[0:0] = [""]
 
 from test import unittest
 from test.pymongo_mocks import DummyMonitor
-from test.utils_shared import AsyncMockPool, parse_read_preference
-from test.utils_shared_selection_tests_shared import (
+from test.utils_selection_tests_shared import (
     get_addresses,
     get_topology_type_name,
     make_server_description,
 )
+from test.utils_shared import parse_read_preference
 
 from bson import json_util
 from pymongo.asynchronous.settings import TopologySettings