PYTHON-5077 - Convert test.test_data_lake to async (#2091)

Author: NoahStapp

test/asynchronous/test_data_lake.py

@@ -0,0 +1,111 @@
+# Copyright 2020-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test Atlas Data Lake."""
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import pytest
+
+sys.path[0:0] = [""]
+
+from test.asynchronous import AsyncIntegrationTest, AsyncUnitTest, async_client_context, unittest
+from test.asynchronous.unified_format import generate_test_classes
+from test.utils import (
+    OvertCommandListener,
+)
+
+from pymongo.asynchronous.helpers import anext
+
+_IS_SYNC = False
+
+pytestmark = pytest.mark.data_lake
+
+
+class TestDataLakeMustConnect(AsyncUnitTest):
+    async def test_connected_to_data_lake(self):
+        data_lake = os.environ.get("TEST_DATA_LAKE")
+        if not data_lake:
+            self.skipTest("TEST_DATA_LAKE is not set")
+
+        self.assertTrue(
+            async_client_context.is_data_lake and async_client_context.connected,
+            "client context must be connected to data lake when DATA_LAKE is set. Failed attempts:\n{}".format(
+                async_client_context.connection_attempt_info()
+            ),
+        )
+
+
+class TestDataLakeProse(AsyncIntegrationTest):
+    # Default test database and collection names.
+    TEST_DB = "test"
+    TEST_COLLECTION = "driverdata"
+
+    @async_client_context.require_data_lake
+    async def asyncSetUp(self):
+        await super().asyncSetUp()
+
+    # Test killCursors
+    async def test_1(self):
+        listener = OvertCommandListener()
+        client = await self.async_rs_or_single_client(event_listeners=[listener])
+        cursor = client[self.TEST_DB][self.TEST_COLLECTION].find({}, batch_size=2)
+        await anext(cursor)
+
+        # find command assertions
+        find_cmd = listener.succeeded_events[-1]
+        self.assertEqual(find_cmd.command_name, "find")
+        cursor_id = find_cmd.reply["cursor"]["id"]
+        cursor_ns = find_cmd.reply["cursor"]["ns"]
+
+        # killCursors command assertions
+        await cursor.close()
+        started = listener.started_events[-1]
+        self.assertEqual(started.command_name, "killCursors")
+        succeeded = listener.succeeded_events[-1]
+        self.assertEqual(succeeded.command_name, "killCursors")
+
+        self.assertIn(cursor_id, started.command["cursors"])
+        target_ns = ".".join([started.command["$db"], started.command["killCursors"]])
+        self.assertEqual(cursor_ns, target_ns)
+
+        self.assertIn(cursor_id, succeeded.reply["cursorsKilled"])
+
+    # Test no auth
+    async def test_2(self):
+        client = await self.async_rs_client_noauth()
+        await client.admin.command("ping")
+
+    # Test with auth
+    async def test_3(self):
+        for mechanism in ["SCRAM-SHA-1", "SCRAM-SHA-256"]:
+            client = await self.async_rs_or_single_client(authMechanism=mechanism)
+            await client[self.TEST_DB][self.TEST_COLLECTION].find_one()
+
+
+# Location of JSON test specifications.
+if _IS_SYNC:
+    TEST_PATH = Path(__file__).parent / "data_lake/unified"
+else:
+    TEST_PATH = Path(__file__).parent.parent / "data_lake/unified"
+
+# Generate unified tests.
+globals().update(generate_test_classes(TEST_PATH, module=__name__))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/test_data_lake.py

@@ -23,20 +23,20 @@
 
 sys.path[0:0] = [""]
 
-from test import IntegrationTest, client_context, unittest
+from test import IntegrationTest, UnitTest, client_context, unittest
 from test.unified_format import generate_test_classes
 from test.utils import (
     OvertCommandListener,
 )
 
-pytestmark = pytest.mark.data_lake
+from pymongo.synchronous.helpers import next
 
+_IS_SYNC = True
 
-# Location of JSON test specifications.
-_TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data_lake")
+pytestmark = pytest.mark.data_lake
 
 
-class TestDataLakeMustConnect(unittest.TestCase):
+class TestDataLakeMustConnect(UnitTest):
     def test_connected_to_data_lake(self):
         data_lake = os.environ.get("TEST_DATA_LAKE")
         if not data_lake:
@@ -55,10 +55,9 @@ class TestDataLakeProse(IntegrationTest):
     TEST_DB = "test"
     TEST_COLLECTION = "driverdata"
 
-    @classmethod
     @client_context.require_data_lake
-    def setUpClass(cls):
-        super().setUpClass()
+    def setUp(self):
+        super().setUp()
 
     # Test killCursors
     def test_1(self):
@@ -99,7 +98,10 @@ def test_3(self):
 
 
 # Location of JSON test specifications.
-TEST_PATH = Path(__file__).parent / "data_lake/unified"
+if _IS_SYNC:
+    TEST_PATH = Path(__file__).parent / "data_lake/unified"
+else:
+    TEST_PATH = Path(__file__).parent.parent / "data_lake/unified"
 
 # Generate unified tests.
 globals().update(generate_test_classes(TEST_PATH, module=__name__))

tools/synchro.py

@@ -203,6 +203,7 @@ def async_only_test(f: str) -> bool:
     "test_crud_unified.py",
     "test_cursor.py",
     "test_database.py",
+    "test_data_lake.py",
     "test_encryption.py",
     "test_grid_file.py",
     "test_logger.py",