(beam) branch master updated: Fix flaky TestDataSampler/GetSamplesForPCollectionsTooManySamples (#38736)

[sh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new cbc2c21edad Fix flaky TestDataSampler/GetSamplesForPCollectionsTooManySamples (#38736)
cbc2c21edad is described below

commit cbc2c21edadf1dc9e7624511b33bb1011b0faffc
Author: Shunping Huang <sh...@google.com>
AuthorDate: Fri May 29 14:52:38 2026 -0400

    Fix flaky TestDataSampler/GetSamplesForPCollectionsTooManySamples (#38736)
    
    Since GetSamples is a destructive operation that clears stored samples,
    polling it in a loop is flaky. If it is called before all asynchronous
    samples have finished processing, it consumes the early samples and clears
    them, causing verifySampledElements to fail.
    
    In this PR, we wait for the sampler to finish before reading to ensure
    all elements are successfully captured.
---
 sdks/go/pkg/beam/core/runtime/exec/datasampler_test.go | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/sdks/go/pkg/beam/core/runtime/exec/datasampler_test.go b/sdks/go/pkg/beam/core/runtime/exec/datasampler_test.go
index 7dce38360af..29ae3a48b5e 100644
--- a/sdks/go/pkg/beam/core/runtime/exec/datasampler_test.go
+++ b/sdks/go/pkg/beam/core/runtime/exec/datasampler_test.go
@@ -104,16 +104,9 @@ func TestDataSampler(t *testing.T) {
 			for _, sample := range test.samples {
 				dataSampler.SendSample(sample.PCollectionID, sample.Element, sample.Timestamp)
 			}
-			var samplesCount = -1
-			var samples map[string][]*DataSample
-			for i := 0; i < 5; i++ {
-				samples = dataSampler.GetSamples(test.pids)
-				if len(samples) == len(test.want) {
-					samplesCount = len(samples)
-					break
-				}
-				time.Sleep(time.Second)
-			}
+			time.Sleep(1 * time.Second)
+			samples := dataSampler.GetSamples(test.pids)
+			samplesCount := len(samples)
 			cancel()
 			if samplesCount != len(test.want) {
 				t.Errorf("got an unexpected number of sampled elements: %v, want: %v", samplesCount, len(test.want))