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