update/add the new prose tests

Author: sleepyStick

test/mockupdb/test_client_metadata.py

@@ -16,7 +16,7 @@
 import time
 import unittest
 from test.utils_shared import CMAPListener
-from typing import Optional
+from typing import Any, Optional
 
 import pytest
 
@@ -43,32 +43,24 @@ class TestClientMetadataProse(unittest.TestCase):
     def setUp(self):
         self.server = MockupDB()
         # there are two handshake requests, i believe one is from the monitor, and the other is from the client
-        self.monitor_handshake = False
         self.handshake_req: Optional[dict] = None
 
         def respond(r):
-            # Only save the very first request from the driver.
-            if self.handshake_req is None:
-                if not self.monitor_handshake:
-                    self.monitor_handshake = True
-                else:
-                    self.handshake_req = r
+            if "ismaster" in r:
+                # then this is a handshake request
+                self.handshake_req = r
             return r.reply(OpMsgReply(minWireVersion=0, maxWireVersion=13))
 
         self.server.autoresponds(respond)
         self.server.run()
         self.addCleanup(self.server.stop)
 
-    def check_metadata_added(
-        self, add_name: str, add_version: Optional[str], add_platform: Optional[str]
-    ) -> None:
-        client = MongoClient(
-            "mongodb://" + self.server.address_string,
-            maxIdleTimeMS=1,
-            driver=DriverInfo("library", "1.2", "Library Platform"),
-        )
-
-        # send initial metadata
+    def send_ping_and_get_metadata(
+        self, client: MongoClient, is_handshake: bool
+    ) -> tuple[str, Optional[str], Optional[str], dict[str, Any]]:
+        # reset
+        if is_handshake:
+            self.handshake_req: Optional[dict] = None
         client.admin.command("ping")
         metadata = _get_handshake_driver_info(self.handshake_req)
         driver_metadata = metadata["driver"]
@@ -77,25 +69,36 @@ def check_metadata_added(
             driver_metadata["version"],
             metadata["platform"],
         )
+        return name, version, platform, metadata
+
+    def check_metadata_added(
+        self,
+        client: MongoClient,
+        add_name: str,
+        add_version: Optional[str],
+        add_platform: Optional[str],
+    ) -> None:
+        # send initial metadata
+        name, version, platform, metadata = self.send_ping_and_get_metadata(client, True)
         time.sleep(0.005)
 
-        # add data
+        # add new metadata
         client._append_metadata(DriverInfo(add_name, add_version, add_platform))
-        # make sure new metadata is being sent
-        self.handshake_req = None
-        client.admin.command("ping")
-        assert self.handshake_req is not None
-        new_metadata = _get_handshake_driver_info(self.handshake_req)
-
-        self.assertEqual(
-            new_metadata["driver"]["name"], f"{name}|{add_name}" if add_name is not None else name
+        new_name, new_version, new_platform, new_metadata = self.send_ping_and_get_metadata(
+            client, True
         )
+        print("IN SEND PING AND GET METADATA")
+        print(name, version, platform)
+        print(metadata)
+        print(new_name, new_version, new_platform)
+        print(new_metadata)
+        self.assertEqual(new_name, f"{name}|{add_name}" if add_name is not None else name)
         self.assertEqual(
-            new_metadata["driver"]["version"],
+            new_version,
             f"{version}|{add_version}" if add_version is not None else version,
         )
         self.assertEqual(
-            new_metadata["platform"],
+            new_platform,
             f"{platform}|{add_platform}" if add_platform is not None else platform,
         )
 
@@ -105,126 +108,100 @@ def check_metadata_added(
         new_metadata.pop("platform")
         self.assertEqual(metadata, new_metadata)
 
-        client.close()
-
     def test_append_metadata(self):
-        self.check_metadata_added("framework", "2.0", "Framework Platform")
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+            driver=DriverInfo("library", "1.2", "Library Platform"),
+        )
+        self.check_metadata_added(client, "framework", "2.0", "Framework Platform")
+        client.close()
 
     def test_append_metadata_platform_none(self):
-        self.check_metadata_added("framework", "2.0", None)
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+            driver=DriverInfo("library", "1.2", "Library Platform"),
+        )
+        self.check_metadata_added(client, "framework", "2.0", None)
+        client.close()
 
     def test_append_metadata_version_none(self):
-        self.check_metadata_added("framework", None, "Framework Platform")
-
-    def test_append_platform_version_none(self):
-        self.check_metadata_added("framework", None, None)
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+            driver=DriverInfo("library", "1.2", "Library Platform"),
+        )
+        self.check_metadata_added(client, "framework", None, "Framework Platform")
+        client.close()
 
-    def test_doesnt_update_established_connections(self):
-        listener = CMAPListener()
+    def test_append_metadata_platform_version_none(self):
         client = MongoClient(
             "mongodb://" + self.server.address_string,
             maxIdleTimeMS=1,
             driver=DriverInfo("library", "1.2", "Library Platform"),
-            event_listeners=[listener],
         )
+        self.check_metadata_added(client, "framework", None, None)
+        client.close()
 
-        # send initial metadata
-        client.admin.command("ping")
-        metadata = _get_handshake_driver_info(self.handshake_req)
-        driver_metadata = metadata["driver"]
-        name, version, platform = (
-            driver_metadata["name"],
-            driver_metadata["version"],
-            metadata["platform"],
+    def test_multiple_successive_metadata_updates(self):
+        client = MongoClient(
+            "mongodb://" + self.server.address_string, maxIdleTimeMS=1, connect=False
         )
-        self.assertIsNotNone(name)
-        self.assertIsNotNone(version)
-        self.assertIsNotNone(platform)
+        client._append_metadata(DriverInfo("library", "1.2", "Library Platform"))
+        self.check_metadata_added(client, "framework", "2.0", "Framework Platform")
+        client.close()
 
-        # add data
-        add_name, add_version, add_platform = "framework", "2.0", "Framework Platform"
-        client._append_metadata(DriverInfo(add_name, add_version, add_platform))
-        # check new data isn't sent
-        self.handshake_req: Optional[dict] = None
-        client.admin.command("ping")
-        self.assertIsNotNone(self.handshake_req)
-        assert self.handshake_req is not None  # so mypy knows that it's not None
-        self.assertNotIn("client", self.handshake_req)
-        self.assertEqual(listener.event_count(ConnectionClosedEvent), 0)
+    def test_multiple_successive_metadata_updates_platform_none(self):
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+        )
+        client._append_metadata(DriverInfo("library", "1.2", "Library Platform"))
+        self.check_metadata_added(client, "framework", "2.0", None)
+        client.close()
 
+    def test_multiple_successive_metadata_updates_version_none(self):
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+        )
+        client._append_metadata(DriverInfo("library", "1.2", "Library Platform"))
+        self.check_metadata_added(client, "framework", None, "Framework Platform")
+        client.close()
+
+    def test_multiple_successive_metadata_updates_platform_version_none(self):
+        client = MongoClient(
+            "mongodb://" + self.server.address_string,
+            maxIdleTimeMS=1,
+        )
+        client._append_metadata(DriverInfo("library", "1.2", "Library Platform"))
+        self.check_metadata_added(client, "framework", None, None)
         client.close()
 
-    def test_append_metadata_multiple_times(self):
+    def test_doesnt_update_established_connections(self):
+        listener = CMAPListener()
         client = MongoClient(
             "mongodb://" + self.server.address_string,
             maxIdleTimeMS=1,
             driver=DriverInfo("library", "1.2", "Library Platform"),
+            event_listeners=[listener],
         )
 
         # send initial metadata
-        client.admin.command("ping")
-        metadata = _get_handshake_driver_info(self.handshake_req)
-        driver_metadata = metadata["driver"]
-        name, version, platform = (
-            driver_metadata["name"],
-            driver_metadata["version"],
-            metadata["platform"],
-        )
-        time.sleep(0.005)
+        name, version, platform, metadata = self.send_ping_and_get_metadata(client, True)
+        self.assertIsNotNone(name)
+        self.assertIsNotNone(version)
+        self.assertIsNotNone(platform)
 
         # add data
         add_name, add_version, add_platform = "framework", "2.0", "Framework Platform"
         client._append_metadata(DriverInfo(add_name, add_version, add_platform))
-        # make sure new metadata is being sent
-        self.handshake_req = None
-        client.admin.command("ping")
-        assert self.handshake_req is not None
-        new_metadata = _get_handshake_driver_info(self.handshake_req)
-
-        self.assertEqual(
-            new_metadata["driver"]["name"], f"{name}|{add_name}" if add_name is not None else name
-        )
-        self.assertEqual(
-            new_metadata["driver"]["version"],
-            f"{version}|{add_version}" if add_version is not None else version,
-        )
-        self.assertEqual(
-            new_metadata["platform"],
-            f"{platform}|{add_platform}" if add_platform is not None else platform,
-        )
-
-        metadata.pop("driver")
-        metadata.pop("platform")
-        new_metadata.pop("driver")
-        new_metadata.pop("platform")
-        self.assertEqual(metadata, new_metadata)
-        time.sleep(0.005)
-
-        # add data again
-        add_name2, add_version2, add_platform2 = "framework2", "3.0", "Framework Platform2"
-        client._append_metadata(DriverInfo(add_name2, add_version2, add_platform2))
-        # make sure new metadata is being sent
-        self.handshake_req = None
+        # check new data isn't sent
+        self.handshake_req: Optional[dict] = None
         client.admin.command("ping")
-        assert self.handshake_req is not None
-        new_metadata2 = _get_handshake_driver_info(self.handshake_req)
-
-        self.assertEqual(
-            new_metadata2["driver"]["name"],
-            f"{name}|{add_name}|{add_name2}" if add_name2 is not None else name,
-        )
-        self.assertEqual(
-            new_metadata2["driver"]["version"],
-            f"{version}|{add_version}|{add_version2}" if add_version2 is not None else version,
-        )
-        self.assertEqual(
-            new_metadata2["platform"],
-            f"{platform}|{add_platform}|{add_platform2}" if add_platform2 is not None else platform,
-        )
-
-        new_metadata2.pop("driver")
-        new_metadata2.pop("platform")
-        self.assertEqual(metadata, new_metadata2)
+        self.assertIsNone(self.handshake_req)
+        self.assertEqual(listener.event_count(ConnectionClosedEvent), 0)
 
         client.close()