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

[mykaul]

Fixes a KeyError in test_idle_heartbeat (tests/integration/standard/test_cluster.py) caused by shard-aware connection replacement racing with the test's connection-identity snapshot.

Background: master already fixed the main source of this via _wait_for_all_shard_connections() (added in 454f727 / f348637), which waits for shard pools to reach their full connection count before the snapshot is taken. That closed the common pool-warm-up race, and deliberately reverted an earlier silent "skip unknown connections" band-aid so the test's assertions stayed meaningful.

That fix doesn't cover a narrower race: HostConnection can still swap an existing connection object for a new one in the same shard slot without changing the pool's connection count — via _open_connection_to_missing_shard() (once orphaned_threshold_reached) or via return_connection() -> _replace() (on a defunct/closed connection). Either can fire during the test's own sleep window, and is more likely under CPU/IO contention (e.g. other tests running concurrently, delaying heartbeat responses).

This PR guards against the resulting KeyError by filtering to connections still present in the snapshot, while keeping the test meaningful:

• In the common case (no replacement), the filtered list equals the full connection list, so behavior/assertions are unchanged from current master.
• If the filter actually drops anything, it now logs a warning so a recurrence stays visible in CI instead of being silently absorbed.
• If every connection was replaced, the test still hard-fails instead of trivially passing.

See PR discussion for more analysis and the connection to the long-standing test_idle_heartbeat flakiness discussed on #653.

Pre-review checklist

• I have split my patch into logically separate commits.
• All commit messages clearly explain what they change and why.
• I added relevant tests for new features and bug fixes.
• All commits compile, pass static checks and pass test.
• PR description sums up the changes and reasons why they should be introduced.
• I have provided docstrings for the public items that I want to introduce.
• I have adjusted the documentation in ./docs/source/.
• I added appropriate Fixes: annotations to PR description.

[mykaul]

CI flakiness - unrelated, will send a separate PR.

[coderabbitai]

Warning

Review limit reached

Next included review available in 28 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: QUIET

Plan: Advanced

Run ID: 3e7bf639-30b9-4802-afc5-6467a92123e4

📥 Commits

Reviewing files that changed from the base of the PR and between 51d1064 and 666bff8.

📒 Files selected for processing (1)

• tests/integration/standard/test_cluster.py

📝 Walkthrough

Walkthrough

The test_idle_heartbeat integration test now tolerates asynchronous connection replacement. It compares snapshot and current connection IDs in both directions, logs discrepancies, and validates only shared connections. The test fails when all snapshot connections disappear. Heartbeat request rotation and idle-state checks now apply only to shared connections.

Possibly related PRs

• scylladb/python-driver#944 — Addresses connection replacement and lifecycle behavior related to this heartbeat test.

Suggested labels: P2, area/Driver_-_python-driver

Suggested reviewers: sylwiaszunejko

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the test, the KeyError, and the connection-replacement fix.
Description check	✅ Passed	The description explains the cause, solution, validation, and checklist status, with only non-critical documentation items unchecked.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mykaul]

Rebased onto current master and dug into this properly. Summary:

What the KeyError actually is: test_idle_heartbeat snapshots each connection's request_ids keyed by id(connection), sleeps for a couple of heartbeat intervals, then looks the connections back up by id() to validate the heartbeats. If a connection object is replaced in between (shard-aware pools can swap connections), the lookup raises KeyError.

Relationship to #653: that PR's review thread flagged this same test's flakiness as a "pre-existing, likely repo-wide intermittent failure," not reproducible locally, possibly triggered by unrelated tests running alongside it — cc @dkropachev @mykaul. Turns out that's exactly right, and traceable to a real mechanism:

• master already fixed the main source, in 454f727 and f348637: wait_for_all_pools() only waits for the first connection per host, so shard-aware connections to the remaining shards could still be opening (and replacing placeholders) during the test's sleep window. _wait_for_all_shard_connections() was added to wait for the connection count to stabilize before the snapshot, and the earlier silent "skip unknown connections" band-aid was deliberately dropped afterward to keep the assertions meaningful.
• That fix guarantees connection count has stabilized, but not that individual connection objects are stable. HostConnection can still swap an existing connection for a new one in the same shard slot, independent of the pool's size, via _open_connection_to_missing_shard() (once orphaned_threshold_reached) or via return_connection() -> _replace() (on a defunct/closed connection). Either can fire mid-test, and is more likely under CPU/IO contention — e.g. other tests running concurrently and delaying heartbeat responses long enough to look orphaned. That lines up well with "why does this pop now, and only when other tests are running alongside it."

So does this fix the #653 root cause? Partially, and it's important to be precise about which part: the primary race (pool warm-up) was already fixed on master independently of this PR. What this PR adds is a second, narrower defense for the residual race described above, which master's fix doesn't cover. I rewrote the fix so it doesn't regress master's "keep assertions meaningful" fix: in the common case (no replacement) the filtered connection list is identical to the full list, so behavior is unchanged; only if the filter actually drops something does it log a warning (visible in CI) and it still hard-fails if every connection was replaced.

Testing:

• python3 -m py_compile tests/integration/standard/test_cluster.py — passes.
• Full integration run isn't possible in my sandbox (no CCM/cluster infra, no network access to resolve/download a Scylla version even for --collect-only); confirmed the same collection error occurs identically against unmodified master, so it's a pre-existing environment limitation, not something this change introduces.
• tests/unit/test_host_connection_pool.py, tests/unit/test_connection.py, tests/unit/test_control_connection.py — 66 passed.
• Full tests/unit/ — 720 passed, 88 skipped (optional reactors unavailable locally), 0 failures.

CI: last recorded run (before rebase) was all green; rebase onto current master will trigger a fresh run.

Squashed into a single commit with an updated message covering all of the above. No unresolved review threads existed on this PR (no reviews yet). Force-pushed the rebase + amended commit; kept as draft per the original intent.

[Copilot]

Pull request overview

Improves test_idle_heartbeat resilience to shard-aware connection replacement races.

Changes:

• Filters validation to snapshotted connections.
• Warns when connections are replaced.
• Fails if no connections remain validatable.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/integration/standard/test_cluster.py`:
- Around line 820-823: Update the assertion message for known_connections so it
accurately states that no snapshotted connection remains in the current
connection list, rather than claiming all connections were replaced. Keep the
existing len(known_connections) > 0 failure condition unchanged.
- Around line 803-810: The reconciliation logic around snapshot_connections must
retain strong references to the original connection objects instead of relying
on integer IDs, preventing ID reuse from misidentifying replacements. Update
missing_from_current, unknown_in_current, and known_connections to compare
connections by identity using is while preserving the existing snapshot and
replacement behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: d1950f42-e109-4890-8d5a-d432f9bc1edf

📥 Commits

Reviewing files that changed from the base of the PR and between 6a7d851 and 51d1064.

📒 Files selected for processing (1)

• tests/integration/standard/test_cluster.py