Fix test_idle_heartbeat: wait for all shard connections before testing

wait_for_all_pools only waits for the first connection per host. Shard-
aware connections to remaining shards are opened asynchronously in the
background. When these background connections complete during the test's
sleep interval, they replace existing connections causing KeyError (old
connection ID not in snapshot) or stale idle state assertions.

Add a helper that polls until all shard connections are established,
then call it after connect(). This stabilizes the connection set before
snapshotting request IDs, eliminating both failure modes and allowing
the original stronger assertions to be restored.

Author: dkropachev

tests/integration/standard/test_cluster.py

@@ -725,12 +725,34 @@ def _warning_are_issued_when_auth(self, auth_provider):
             assert auth_warning >= 4
             assert auth_warning == mock_handler.get_message_count("debug", "Got ReadyMessage on new connection")
 
+    def _wait_for_all_shard_connections(self, cluster, timeout=30):
+        """Wait until all shard-aware connections are fully established."""
+        from cassandra.pool import HostConnection
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            all_connected = True
+            for holder in cluster.get_connection_holders():
+                if not isinstance(holder, HostConnection):
+                    continue
+                if holder.host.sharding_info and len(holder._connections) < holder.host.sharding_info.shards_count:
+                    all_connected = False
+                    break
+            if all_connected:
+                return
+            time.sleep(0.1)
+        raise RuntimeError("Timed out waiting for all shard connections to be established")
+
     def test_idle_heartbeat(self):
         interval = 2
         cluster = TestCluster(idle_heartbeat_interval=interval,
                               monitor_reporting_enabled=False)
         session = cluster.connect(wait_for_all_pools=True)
 
+        # wait_for_all_pools only waits for the first connection per host;
+        # shard-aware connections to remaining shards are opened in background.
+        # Wait for them to stabilize so they don't get replaced during the test.
+        self._wait_for_all_shard_connections(cluster)
+
         # This test relies on impl details of connection req id management to see if heartbeats
         # are being sent. May need update if impl is changed
         connection_request_ids = {}
@@ -746,11 +768,8 @@ def test_idle_heartbeat(self):
 
         connections = [c for holders in cluster.get_connection_holders() for c in holders.get_connections()]
 
-        # make sure requests were sent on all connections that existed before the sleep
-        # (shard-aware reconnection may replace connections during the sleep interval)
+        # make sure requests were sent on all connections
         for c in connections:
-            if id(c) not in connection_request_ids:
-                continue
             expected_ids = connection_request_ids[id(c)]
             expected_ids.rotate(-1)
             with c.lock: