commit 8226412ffec59d05903f1c820d7c052eaaa4f1c2
author Daniel Thevessen <danthev@google.com> Fri Sep 17 16:30:08 2021 -0700
committer Daniel Thevessen <danthev@google.com> Fri Sep 17 16:30:08 2021 -0700
tree df145130ce964d76efaface9fc6573afadd000f8
parent 8072cc0bcfd4eee08a95902e13b9bf1dc2338693

Fix flaky test.

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..51e4d3e 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,17 @@
     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;
   }