From 9dd2e72828c13bf088403ed789c4672ffef99564 Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Wed, 15 Jul 2026 02:20:34 -0500 Subject: [PATCH] fix flaky multi example connector consistency (#48206) Signed-off-by: aarushjain29 Co-authored-by: Cursor Co-authored-by: Andreas Karatzas --- .../kv_connector/unit/test_multi_connector.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 441e8e3676..14ea6759f4 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -318,8 +318,12 @@ def test_multi_example_connector_consistency(): # connector (first nonzero match is chosen), so update_state_after_alloc # will report those external tokens on that one. Other connectors still # receive the request's real blocks but with 0 external tokens. - storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) - storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) + storage1_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage1-SCHEDULER"]) + ) + storage2_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage2-SCHEDULER"]) + ) assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", @@ -348,8 +352,12 @@ def test_multi_example_connector_consistency(): # return 0 from the first connector, while the second connector has a hit. # Both connectors receive the request's real blocks, but only the chosen # (second) connector reports external tokens. - storage1_scheduler_events = _ignore_event_collection(events["storage1-SCHEDULER"]) - storage2_scheduler_events = _ignore_event_collection(events["storage2-SCHEDULER"]) + storage1_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage1-SCHEDULER"]) + ) + storage2_scheduler_events = _events_from_request( + _ignore_event_collection(events["storage2-SCHEDULER"]) + ) assert storage1_scheduler_events[:4] == [ "on_new_request", "get_num_new_matched_tokens 0", @@ -375,6 +383,16 @@ def _ignore_event_collection(events: list[str]) -> list[str]: return [event for event in events if event not in ignored] +def _events_from_request(events: list[str]) -> list[str]: + # The async engine-core can emit a trailing build_connector_meta from the + # previous generation's final (idle) scheduler step. Depending on timing it + # may be flushed into this window ahead of on_new_request, so anchor the + # comparison on the new request's first event to avoid a flaky ordering. + if "on_new_request" in events: + return events[events.index("on_new_request") :] + return events + + def get_connector_events() -> dict[str, list[str]]: # Read in connector events and reset the files. import glob