PYTHON-5366 - test_pool_reset waits until Pool.reset() increments gen… (#2797)

Co-authored-by: Qi Deng <qdeng@aurascape.ai>

pymongo/_azure_helpers.py

@@ -17,6 +17,7 @@ from __future__ import annotations
 
 import json
 from typing import Any, Optional
+from urllib.parse import quote
 
 
 def _get_azure_response(
@@ -29,7 +30,7 @@ def _get_azure_response(
     url += "?api-version=2018-02-01"
     url += f"&resource={resource}"
     if client_id:
-        url += f"&client_id={client_id}"
+        url += f"&client_id={quote(client_id)}"
     headers = {"Metadata": "true", "Accept": "application/json"}
     request = Request(url, headers=headers)  # noqa: S310
     try:

test/asynchronous/test_client.py

@@ -2711,11 +2711,11 @@ class TestClientPool(AsyncMockClientTest):
 
         await async_wait_until(lambda: len(c.nodes) == 1, "connect")
         self.assertEqual(await c.address, ("c", 3))
-        # Assert that we create 1 pooled connection.
+        # Wait for the pooled connection to be registered
         await listener.async_wait_for_event(monitoring.ConnectionReadyEvent, 1)
         self.assertEqual(listener.event_count(monitoring.ConnectionCreatedEvent), 1)
         arbiter = c._topology.get_server_by_address(("c", 3))
-        self.assertEqual(len(arbiter.pool.conns), 1)
+        await async_wait_until(lambda: len(arbiter.pool.conns) == 1, "create 1 pooled connection")
         # Arbiter pool is marked ready.
         self.assertEqual(listener.event_count(monitoring.PoolReadyEvent), 1)
 

test/test_azure_helpers.py

@@ -150,6 +150,20 @@ class TestGetAzureResponse(unittest.TestCase):
         _, kwargs = mock_open.call_args
         self.assertEqual(kwargs["timeout"], 42)
 
+    def test_client_id_is_url_encoded(self):
+        """Ensure special characters in client_id are percent-encoded."""
+        body = json.dumps({"access_token": "tok", "expires_in": "3600"})
+        with _mock_urlopen(200, body) as mock_open:
+            self._call(client_id="id with spaces&special=chars")
+
+        url = mock_open.call_args[0][0].full_url
+        # '&' and '=' must be percent-encoded so they don't inject extra query params
+        self.assertIn("client_id=id%20with%20spaces%26special%3Dchars", url)
+        # The encoded client_id should not introduce a raw '&'
+        # Count params: api-version, resource, client_id — exactly 3
+        query_string = url.split("?", 1)[1]
+        self.assertEqual(query_string.count("&"), 2)
+
 
 if __name__ == "__main__":
     unittest.main()

test/test_client.py

@@ -2666,11 +2666,11 @@ class TestClientPool(MockClientTest):
 
         wait_until(lambda: len(c.nodes) == 1, "connect")
         self.assertEqual(c.address, ("c", 3))
-        # Assert that we create 1 pooled connection.
+        # Wait for the pooled connection to be registered
         listener.wait_for_event(monitoring.ConnectionReadyEvent, 1)
         self.assertEqual(listener.event_count(monitoring.ConnectionCreatedEvent), 1)
         arbiter = c._topology.get_server_by_address(("c", 3))
-        self.assertEqual(len(arbiter.pool.conns), 1)
+        wait_until(lambda: len(arbiter.pool.conns) == 1, "create 1 pooled connection")
         # Arbiter pool is marked ready.
         self.assertEqual(listener.event_count(monitoring.PoolReadyEvent), 1)
 

test/test_topology.py

@@ -23,7 +23,7 @@ sys.path[0:0] = [""]
 
 from test import client_knobs, unittest
 from test.pymongo_mocks import DummyMonitor
-from test.utils import MockPool, flaky
+from test.utils import MockPool
 from test.utils_shared import wait_until
 
 from bson.objectid import ObjectId
@@ -755,7 +755,6 @@ def wait_for_primary(topology):
 class TestTopologyErrors(TopologyTest):
     # Errors when calling hello.
 
-    @flaky(reason="PYTHON-5366")
     def test_pool_reset(self):
         # hello succeeds at first, then always raises socket error.
         hello_count = [0]
@@ -776,7 +775,11 @@ class TestTopologyErrors(TopologyTest):
 
         # Pool is reset by hello failure.
         t.request_check_all()
-        self.assertNotEqual(generation, server.pool.gen.get_overall())
+        # Wait for the monitor's hello failure to trigger Pool.reset() and bump the generation.
+        wait_until(
+            lambda: server.pool.gen.get_overall() != generation,
+            "pool reset after failed monitor check",
+        )
 
     def test_hello_retry(self):
         # hello succeeds at first, then raises socket error, then succeeds.