(fix) test: Fix KeyError in test_idle_heartbeat with connection repla…

tests/integration/standard/test_cluster.py

@@ -776,22 +776,64 @@ def test_idle_heartbeat(self):
                 # make sure none are idle (should have startup messages
                 assert not c.is_idle
                 with c.lock:
-                    connection_request_ids[id(c)] = deque(c.request_ids)  # copy of request ids
+                    # keyed by the connection object itself (not id(c)) so the
+                    # reference stays alive and a freed id can't be reused
+                    connection_request_ids[c] = deque(c.request_ids)  # copy of request ids
 
         # let two heatbeat intervals pass (first one had startup messages in it)
         time.sleep(2 * interval + interval/2)
 
         connections = [c for holders in cluster.get_connection_holders() for c in holders.get_connections()]
 
-        # make sure requests were sent on all connections
-        for c in connections:
-            expected_ids = connection_request_ids[id(c)]
+        # _wait_for_all_shard_connections() above prevents the common KeyError caused by
+        # shard connections still being opened during pool warm-up. It does not cover a
+        # narrower, still-possible race: HostConnection can swap an existing connection
+        # for a new object in the same shard slot independently of the connection count,
+        # e.g. via _open_connection_to_missing_shard() once orphaned_threshold_reached, or
+        # via return_connection()/_replace() when a connection is defunct/closed. The
+        # latter has a genuine asynchronous gap: return_connection() pops the old
+        # connection out of the pool's dict synchronously, but _replace() (which inserts
+        # the replacement) runs later on another thread via session.submit(). A snapshot
+        # taken during that gap sees fewer *current* connections than were recorded, and
+        # every one of them is still "known" (nothing new has been inserted yet) - so
+        # checking only "is every current connection known" can't distinguish "fewer
+        # connections than snapshotted" from "same connections, none replaced". Compare
+        # both directions of the snapshot/current intersection so partial removals (as
+        # well as outright replacements) stay visible, and log loudly if either side
+        # actually differs so a recurrence stays visible in CI instead of being silently
+        # absorbed.
+        current_conns = set(connections)
+        snapshot_conns = set(connection_request_ids)
+        # snapshotted connections no longer present at all (removed, replacement may
+        # not have landed yet)
+        missing_from_current = snapshot_conns - current_conns
+        # connections present now that weren't in the snapshot (replacement landed)
+        unknown_in_current = current_conns - snapshot_conns
+        known_connections = [c for c in connections if c in connection_request_ids]
+        if missing_from_current or unknown_in_current:
+            log.warning(
+                "test_idle_heartbeat: connections changed between the snapshot and "
+                "validation (%d snapshotted connections missing, %d new connections "
+                "observed, out of %d snapshotted / %d current); validating only the "
+                "%d connections common to both",
+                len(missing_from_current), len(unknown_in_current),
+                len(snapshot_conns), len(current_conns), len(known_connections)
+            )
+        assert len(known_connections) > 0, (
+            "No snapshotted connections remained in the current pools; "
+            "no heartbeats could be validated"
+        )
+
+        # make sure heartbeat requests were sent on all known connections
+        for c in known_connections:
+            expected_ids = connection_request_ids[c]
             expected_ids.rotate(-1)
             with c.lock:
                 assertListEqual(list(c.request_ids), list(expected_ids))
 
-        # assert idle status
-        assert all(c.is_idle for c in connections)
+        # assert idle status on known connections only (replaced connections
+        # may not have had their idle state set yet)
+        assert all(c.is_idle for c in known_connections)
 
         # send enough messages to ensure all connections are used
         # (with shard-aware routing, each query only hits one shard per host,