Use a mock random number generator to make the unit test flaky-proof

RokLenarcic · Jan 22, 2019 · 86cad39 · 86cad39
1 parent a278504
commit 86cad39
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 7 deletions.
diff --git a/zookeeper-server/src/main/java/org/apache/zookeeper/server/BlueThrottle.java b/zookeeper-server/src/main/java/org/apache/zookeeper/server/BlueThrottle.java
@@ -215,7 +215,7 @@ else if ((length <= threshold) && (drop > 0)) {
             lastFreeze = now;
         }
 
-        if (rng.nextDouble() <= drop) {
+        if (rng.nextDouble() < drop) {
             return false;
         }
         return true;

diff --git a/zookeeper-server/src/test/java/org/apache/zookeeper/server/BlueThrottleTest.java b/zookeeper-server/src/test/java/org/apache/zookeeper/server/BlueThrottleTest.java
@@ -19,13 +19,42 @@
 
 import org.apache.zookeeper.ZKTestCase;
 import org.junit.Assert;
+import org.junit.Rule;
 import org.junit.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Random;
+
 public class BlueThrottleTest extends ZKTestCase {
     private static final Logger LOG = LoggerFactory.getLogger(BlueThrottleTest.class);
 
+    class MockRandom extends Random {
+        int flag = 0;
+        BlueThrottle throttle;
+
+        public MockRandom(BlueThrottle bt) {
+            this.throttle = bt;
+        }
+
+        @Override
+        public double nextDouble() {
+            if (throttle.getDropChance() > 0) {
+                flag = 1 - flag;
+                return flag;
+            } else {
+                return 1;
+            }
+        }
+    }
+
+    class BlueThrottleWithMockRandom extends BlueThrottle {
+        public BlueThrottleWithMockRandom() {
+            super();
+            this.rng = new MockRandom(this);
+        }
+    }
+
     @Test
     public void testThrottleDisabled() {
         BlueThrottle throttler = new BlueThrottle();
@@ -56,8 +85,8 @@ public void testThrottleWithRefill() throws InterruptedException {
 
     @Test
     public void testThrottleWithoutRandomDropping() throws InterruptedException {
-        int maxTokens = 10;
-        BlueThrottle throttler = new BlueThrottle();
+        int maxTokens = 5;
+        BlueThrottle throttler = new BlueThrottleWithMockRandom();
         throttler.setMaxTokens(maxTokens);
         throttler.setFillCount(maxTokens);
         throttler.setFillTime(1000);
@@ -85,8 +114,8 @@ public void testThrottleWithoutRandomDropping() throws InterruptedException {
 
     @Test
     public void testThrottleWithRandomDropping() throws InterruptedException {
-        int maxTokens = 10;
-        BlueThrottle throttler = new BlueThrottle();
+        int maxTokens = 5;
+        BlueThrottle throttler = new BlueThrottleWithMockRandom();
         throttler.setMaxTokens(maxTokens);
         throttler.setFillCount(maxTokens);
         throttler.setFillTime(1000);
@@ -97,14 +126,14 @@ public void testThrottleWithRandomDropping() throws InterruptedException {
             throttler.checkLimit(1);
         Assert.assertEquals("All tokens should be used up by now", throttler.getMaxTokens(), throttler.getDeficit());
 
-        Thread.sleep(110);
+        Thread.sleep(120);
         //this will trigger dropping probability being increased
         throttler.checkLimit(1);
         Assert.assertTrue("Dropping probability should be increased", throttler.getDropChance()>0);
         LOG.info("Dropping probability is {}", throttler.getDropChance());
 
         //allow bucket to be refilled
-        Thread.sleep(1500);
+        Thread.sleep(1100);
         LOG.info("Bucket is refilled with {} tokens.", maxTokens);
 
         int accepted = 0;