test: Make EmbeddedKafkaSupervisorTest wait for Broker queryability

[amaechler]

Description

EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions (added in #20474) sometimes fails its first SQL check with Object 'datasource_xxx' not found. It failed all 4 attempts in one downstream CI job. Two other tests in the class have the same pattern, so this PR fixes them too.

Root cause

The test waited for the Indexer's ingest/events/processed metric and then queried the Broker. That metric only means the task read the row. The Broker may not have added the datasource to its SQL schema yet.

The second half of the test has a related race, which Claude found by reading the code and I have not reproduced. Replacing the supervisor spec stops the old task without publishing. If its partition 1 offset was not committed yet, the new task reads partition 1 again, and that alone can satisfy the cumulative wait before partition 3 arrives.

Fix

• Add EmbeddedClusterApis.waitForSegmentsToBeQueryable(dataSource, numSegments, timeoutMillis). It works like waitForAllSegmentsToBeQueryable but takes the segment count, so it also covers realtime segments that are not published yet. waitForAllSegmentsToBeQueryable now calls it.
• testSupervisorIngestsOnlySelectedPartitions and test_runKafkaSupervisorWithHeaderFiltering share a waitForHandedOffRows helper: wait for handoff, then for the Broker to see the segments. Handoff happens after publish, so the offsets are committed before the spec is replaced. The selected-partitions test also uses submitSupervisor now, so the @AfterEach teardown cleans up after a failure.
• test_runKafkaSupervisor checks the row count while ingestion is still running, so it waits for the Broker to see all 10 segments instead of waiting for handoff.

Related

Earlier fixes in this class: #19581 (the events/processed wait this PR replaces), #20479 (the @AfterEach teardown) and #20519 (waiting for handoff without suspending the supervisor).

---

This PR has:

• been self-reviewed.
• added comments explaining the "why" and the intent of the code wherever would not be obvious for an unfamiliar reader.
• added unit tests or modified existing tests to cover new code paths, ensuring the threshold for code coverage is met.

[FrankChen021]

The daily flaky test triage does not report the failure of this test, could you show me a CI job link so that I can investigate why it was not reported?

[amaechler]

Thanks for checking this out Frank. I can't share the failure with you here because it happened in CI on our private downstream fork. The test file is the same as upstream though. Here is the relevant part of the log; all 4 surefire attempts failed at the same line:

[ERROR] Tests run: 5, Failures: 0, Errors: 1, Skipped: 0, Time elapsed: 39.55 s <<< FAILURE! -- in org.apache.druid.indexing.kafka.simulate.EmbeddedKafkaSupervisorTest
[ERROR] org.apache.druid.indexing.kafka.simulate.EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions -- Time elapsed: 0.569 s <<< ERROR!
java.lang.RuntimeException: java.lang.RuntimeException: org.apache.druid.rpc.HttpResponseException: Server error [400 Bad Request]; body: {"error":"druidException","errorCode":"invalidInput",...,"errorMessage":"Object 'datasource_mfoojeni' not found (line [1], column [22])",...}
	at org.apache.druid.testing.embedded.EmbeddedServiceClient.runSql(EmbeddedServiceClient.java:252)
	at org.apache.druid.testing.embedded.EmbeddedClusterApis.runSql(EmbeddedClusterApis.java:150)
	at org.apache.druid.testing.embedded.EmbeddedDruidCluster.runSql(EmbeddedDruidCluster.java:360)
	at org.apache.druid.indexing.kafka.simulate.EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions(EmbeddedKafkaSupervisorTest.java:130)
...
[ERROR]   Run 1: EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions:130 » HttpResponseException 400: Object 'datasource_mfoojeni' not found (line [1], column [22])
[ERROR]   Run 2: EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions:130 » HttpResponseException 400: Object 'datasource_bgbbclne' not found (line [1], column [22])
[ERROR]   Run 3: EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions:130 » HttpResponseException 400: Object 'datasource_bedaldci' not found (line [1], column [22])
[ERROR]   Run 4: EmbeddedKafkaSupervisorTest.testSupervisorIngestsOnlySelectedPartitions:130 » HttpResponseException 400: Object 'datasource_mejaofbp' not found (line [1], column [22])

Claude was also able to reproduce it on apache master: with -Dsurefire.rerunFailingTestsCount=0 and a temporary @RepeatedTest(10) on the method, 3 of 10 runs failed with the same error at line 130 if you want to give it a shot.

[FrankChen021]

Noted and thanks for the fix.

just curious that the test didn't fail for commits on master branch and other PRs after that one was merged.
Instead, #20604 caused lots of failures.

Tests of ingestion are very like flaky based on the triage report, I haven't come up with ideas of how to make sure they are stable before merging.

[FrankChen021]

LGTM. This matches both ways testSupervisorIngestsOnlySelectedPartitions failed on master on 2026-10-16. Both runs passed on retry, so neither job failed:

• Line 130, ac864bf (#20583), job 116764398036: the first SELECT COUNT(*) returned 400 invalidInput: Object 'datasource_fpedfoba' not found. That is the Broker-schema race in the description: ingest/events/processed fired on the Indexer before the Broker knew about the datasource.
• Line 140, 17d5ad0 (#20584), job 116764225473: expected: <2> but was: <1> after replacing the spec with partitions {1, 3}. This is CI evidence for the second race the description calls unreproduced. The cumulative events/processed >= 2 can pass on a re-read of p1 (its offset not yet committed) or before the Broker sees the p3 segment. Waiting for handoff and then Broker queryability closes both.

What I checked:

• waitForHandedOffRows will finish: newKafkaSupervisor() uses maxRowsPerSegment=1 and taskDuration=500ms, so each row is its own segment and publishes quickly. dataSource is unique per test and there is one task group with no replicas, so hasSumAtLeast(2) can only count p1 and p3.
• Handoff follows a transactional publish, so p1's offset is committed before the spec is replaced and the new task can't re-read it.
• submitSupervisor updates supervisorSpec, so @AfterEach suspends the expanded spec. Removing the trailing createSuspendedSpec() is safe, and cleanup now also runs when an assertion fails.
• The EmbeddedClusterApis refactor keeps waitForAllSegmentsToBeQueryable's behaviour. The >= in waitForSegmentsToBeQueryable is right for realtime segments that aren't in the metadata store yet.
• CI is green on ad39f699, including the E*,G*,F* and K* jobs.
• Locally, EmbeddedKafkaSupervisorTest passed 5 of 5 tests in each of 5 back-to-back runs on ad39f699 with -Dsurefire.rerunFailingTestsCount=0 (JDK 25, Colima, about 24–25 s per run).

Nit, not blocking: the description could link the line-140 failure above as evidence for the second race.

Review prepared with Claude Code (AI-assisted); I checked the findings before approving.