Fix test import

NoahStapp · NoahStapp · commit 81ea0bcb2bf4 · 2024-10-16T09:18:16.000-04:00
diff --git a/test/conftest.py b/test/conftest.py
@@ -2,8 +2,7 @@
 
 import asyncio
 import sys
-from test import pytest_conf
-from test.synchronous import setup, teardown
+from test import pytest_conf, setup, teardown
 
 import pytest
 
diff --git a/test/pymongo_mocks.py b/test/pymongo_mocks.py
@@ -19,7 +19,6 @@
 import weakref
 from functools import partial
 from test import client_context
-from test.synchronous import client_context
 
 from pymongo import MongoClient, common
 from pymongo.errors import AutoReconnect, NetworkTimeout
diff --git a/test/test_auth.py b/test/test_auth.py
@@ -23,7 +23,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     PyMongoTestCase,
     SkipTest,
diff --git a/test/test_auth_spec.py b/test/test_auth_spec.py
@@ -20,7 +20,7 @@
 import os
 import sys
 import warnings
-from test.synchronous import PyMongoTestCase
+from test import PyMongoTestCase
 
 sys.path[0:0] = [""]
 
diff --git a/test/test_bulk.py b/test/test_bulk.py
@@ -23,7 +23,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import IntegrationTest, client_context, remove_all_users, unittest
+from test import IntegrationTest, client_context, remove_all_users, unittest
 from test.utils import wait_until
 
 from bson.binary import Binary, UuidRepresentation
diff --git a/test/test_change_stream.py b/test/test_change_stream.py
@@ -28,7 +28,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     PyMongoTestCase,
     Version,
diff --git a/test/test_client.py b/test/test_client.py
@@ -44,7 +44,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     HAVE_IPADDRESS,
     IntegrationTest,
     MockClientTest,
@@ -58,7 +58,7 @@
     remove_all_users,
     unittest,
 )
-from test.synchronous.pymongo_mocks import MockClient
+from test.pymongo_mocks import MockClient
 from test.test_binary import BinaryData
 from test.utils import (
     NTHREADS,
diff --git a/test/test_client_bulk_write.py b/test/test_client_bulk_write.py
@@ -20,7 +20,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     client_context,
     unittest,
diff --git a/test/test_client_context.py b/test/test_client_context.py
@@ -18,7 +18,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import SkipTest, UnitTest, client_context, unittest
+from test import SkipTest, UnitTest, client_context, unittest
 
 _IS_SYNC = True
 
diff --git a/test/test_collation.py b/test/test_collation.py
@@ -17,7 +17,7 @@
 
 import functools
 import warnings
-from test.synchronous import IntegrationTest, client_context, unittest
+from test import IntegrationTest, client_context, unittest
 from test.utils import EventListener
 from typing import Any
 
diff --git a/test/test_collection.py b/test/test_collection.py
@@ -27,11 +27,11 @@
 
 sys.path[0:0] = [""]
 
-from test import unittest
-from test.synchronous import (  # TODO: fix sync imports in PYTHON-4528
+from test import (  # TODO: fix sync imports in PYTHON-4528
     IntegrationTest,
     UnitTest,
     client_context,
+    unittest,
 )
 from test.utils import (
     IMPOSSIBLE_WRITE_CONCERN,
diff --git a/test/test_common.py b/test/test_common.py
@@ -20,7 +20,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import IntegrationTest, client_context, connected, unittest
+from test import IntegrationTest, client_context, connected, unittest
 
 from bson.binary import PYTHON_LEGACY, STANDARD, Binary, UuidRepresentation
 from bson.codec_options import CodecOptions
diff --git a/test/test_connections_survive_primary_stepdown_spec.py b/test/test_connections_survive_primary_stepdown_spec.py
@@ -19,13 +19,13 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     client_context,
     reset_client_context,
     unittest,
 )
-from test.synchronous.helpers import repl_set_step_down
+from test.helpers import repl_set_step_down
 from test.utils import (
     CMAPListener,
     ensure_all_connected,
diff --git a/test/test_cursor.py b/test/test_cursor.py
@@ -29,7 +29,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import IntegrationTest, client_context, unittest
+from test import IntegrationTest, client_context, unittest
 from test.utils import (
     AllowListEventListener,
     EventListener,
diff --git a/test/test_database.py b/test/test_database.py
@@ -23,8 +23,7 @@
 
 sys.path[0:0] = [""]
 
-from test import unittest
-from test.synchronous import IntegrationTest, client_context
+from test import IntegrationTest, client_context, unittest
 from test.test_custom_types import DECIMAL_CODECOPTS
 from test.utils import (
     IMPOSSIBLE_WRITE_CONCERN,
diff --git a/test/test_encryption.py b/test/test_encryption.py
@@ -28,8 +28,8 @@
 import traceback
 import uuid
 import warnings
-from test.synchronous import IntegrationTest, PyMongoTestCase, client_context
-from test.synchronous.test_bulk import BulkTestBase
+from test import IntegrationTest, PyMongoTestCase, client_context
+from test.test_bulk import BulkTestBase
 from threading import Thread
 from typing import Any, Dict, Mapping, Optional
 
@@ -53,8 +53,7 @@
     KMIP_CREDS,
     LOCAL_MASTER_KEY,
 )
-from test.synchronous.test_bulk import BulkTestBase
-from test.synchronous.utils_spec_runner import SpecRunner
+from test.test_bulk import BulkTestBase
 from test.unified_format import generate_test_classes
 from test.utils import (
     AllowListEventListener,
@@ -65,6 +64,7 @@
     is_greenthread_patched,
     wait_until,
 )
+from test.utils_spec_runner import SpecRunner
 
 from bson import DatetimeMS, Decimal128, encode, json_util
 from bson.binary import UUID_SUBTYPE, Binary, UuidRepresentation
diff --git a/test/test_grid_file.py b/test/test_grid_file.py
@@ -21,7 +21,7 @@
 import sys
 import zipfile
 from io import BytesIO
-from test.synchronous import (
+from test import (
     IntegrationTest,
     UnitTest,
     client_context,
diff --git a/test/test_logger.py b/test/test_logger.py
@@ -14,8 +14,7 @@
 from __future__ import annotations
 
 import os
-from test import unittest
-from test.synchronous import IntegrationTest
+from test import IntegrationTest, unittest
 from unittest.mock import patch
 
 from bson import json_util
diff --git a/test/test_monitoring.py b/test/test_monitoring.py
@@ -22,7 +22,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     client_context,
     client_knobs,
diff --git a/test/test_raw_bson.py b/test/test_raw_bson.py
@@ -19,7 +19,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import IntegrationTest, client_context, unittest
+from test import IntegrationTest, client_context, unittest
 
 from bson import Code, DBRef, decode, encode
 from bson.binary import JAVA_LEGACY, Binary, UuidRepresentation
diff --git a/test/test_retryable_reads.py b/test/test_retryable_reads.py
@@ -24,7 +24,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     PyMongoTestCase,
     client_context,
diff --git a/test/test_retryable_writes.py b/test/test_retryable_writes.py
@@ -23,13 +23,13 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     SkipTest,
     client_context,
     unittest,
 )
-from test.synchronous.helpers import client_knobs
+from test.helpers import client_knobs
 from test.utils import (
     CMAPListener,
     DeprecationFilter,
diff --git a/test/test_session.py b/test/test_session.py
@@ -25,7 +25,7 @@
 
 sys.path[0:0] = [""]
 
-from test.synchronous import (
+from test import (
     IntegrationTest,
     PyMongoTestCase,
     SkipTest,
diff --git a/test/test_transactions.py b/test/test_transactions.py
@@ -17,13 +17,13 @@
 
 import sys
 from io import BytesIO
-from test.synchronous.utils_spec_runner import SpecRunner
+from test.utils_spec_runner import SpecRunner
 
 from gridfs.synchronous.grid_file import GridFS, GridFSBucket
 
 sys.path[0:0] = [""]
 
-from test.synchronous import IntegrationTest, client_context, unittest
+from test import IntegrationTest, client_context, unittest
 from test.utils import (
     OvertCommandListener,
     wait_until,
diff --git a/test/utils_spec_runner.py b/test/utils_spec_runner.py
@@ -18,7 +18,7 @@
 import functools
 import threading
 from collections import abc
-from test.synchronous import IntegrationTest, client_context, client_knobs
+from test import IntegrationTest, client_context, client_knobs
 from test.utils import (
     CMAPListener,
     CompareType,
diff --git a/tools/synchro.py b/tools/synchro.py
@@ -228,6 +228,8 @@ def process_files(files: list[str]) -> None:
                 lines = translate_async_sleeps(lines)
                 if file in docstring_translate_files:
                     lines = translate_docstrings(lines)
+                if file in sync_test_files:
+                    translate_imports(lines)
                 f.seek(0)
                 f.writelines(lines)
                 f.truncate()
@@ -259,6 +261,15 @@ def translate_coroutine_types(lines: list[str]) -> list[str]:
     return lines
 
 
+def translate_imports(lines: list[str]) -> list[str]:
+    for k, v in import_replacements.items():
+        matches = [line for line in lines if k in line and "import" in line]
+        for line in matches:
+            index = lines.index(line)
+            lines[index] = line.replace(k, v)
+    return lines
+
+
 def translate_async_sleeps(lines: list[str]) -> list[str]:
     blocking_sleeps = [line for line in lines if "asyncio.sleep(0)" in line]
     lines = [line for line in lines if line not in blocking_sleeps]
