From 38f5f58435de3e140dda2c72f1124bb8f910d93d Mon Sep 17 00:00:00 2001
From: Sam Whittle <samuelw@google.com>
Date: Tue, 18 Feb 2025 13:38:12 +0000
Subject: [PATCH] Fix race in FanOutStreamingEngineWorkerHarnessTest

Multiple assignments can be racily collapsed causing the verification
to block.
---
 .../harness/FanOutStreamingEngineWorkerHarnessTest.java         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
index e2815e3a54de..9073a0230c4f 100644
--- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
+++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarnessTest.java
@@ -325,8 +325,8 @@ public void testOnNewWorkerMetadata_redistributesBudget() throws InterruptedExce
             noOpProcessWorkItemFn());
 
     fakeGetWorkerMetadataStub.injectWorkerMetadata(firstWorkerMetadata);
+    verify(getWorkBudgetDistributor, times(1)).distributeBudget(any(), any());
     fakeGetWorkerMetadataStub.injectWorkerMetadata(secondWorkerMetadata);
-
     verify(getWorkBudgetDistributor, times(2)).distributeBudget(any(), any());
   }