commit a7b706cb9d1e84709f89fe98d1dda94d4eb1243b
author tvalentyn <tvalentyn@users.noreply.github.com> Mon Sep 20 10:08:39 2021 -0700
committer GitHub <noreply@github.com> Mon Sep 20 10:08:39 2021 -0700
tree 7be418ac0ab3da0142b3a43e1f974b48cc542450
parent 8072cc0bcfd4eee08a95902e13b9bf1dc2338693
parent 31b7056353522b130d0c0123fb5d165f57684b82

[BEAM-12858] Fix flaky test #15534

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java

diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
index 6c57a84..604b5cb 100644
--- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
+++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/datastore/RampupThrottlingFnTest.java
@@ -52,11 +52,18 @@
     MockitoAnnotations.openMocks(this);
 
     DateTimeUtils.setCurrentMillisFixed(0);
-    RampupThrottlingFn<Void> rampupThrottlingFn = new RampupThrottlingFn<>(1);
+    RampupThrottlingFn<Void> rampupThrottlingFn =
+        new RampupThrottlingFn<Void>(1) {
+          @Override
+          @Setup
+          public void setup() {
+            super.setup();
+            this.sleeper = mockSleeper;
+          }
+        };
     rampupThrottlingFnTester = DoFnTester.of(rampupThrottlingFn);
     rampupThrottlingFnTester.setCloningBehavior(CloningBehavior.DO_NOT_CLONE);
     rampupThrottlingFnTester.startBundle();
-    rampupThrottlingFn.sleeper = mockSleeper;
     rampupThrottlingFn.throttlingMsecs = mockCounter;
   }