From af1b185ce461d5826691243b41e95d3ab0f934bf Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Tue, 3 Apr 2018 16:51:57 -0700
Subject: [PATCH] Fix flaky deadline tests

---
 nomad/drainer/drain_heap_test.go | 17 +++++++++--------
 testutil/wait.go                 |  5 +++++
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/nomad/drainer/drain_heap_test.go b/nomad/drainer/drain_heap_test.go
index 464c413a2..a6574fc16 100644
--- a/nomad/drainer/drain_heap_test.go
+++ b/nomad/drainer/drain_heap_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/hashicorp/nomad/testutil"
 	"github.com/stretchr/testify/require"
 )
 
@@ -29,7 +30,7 @@ func TestDeadlineHeap_WatchAndGet(t *testing.T) {
 	var batch []string
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(2 * wait):
+	case <-time.After(testutil.Timeout(2 * wait)):
 		t.Fatal("timeout")
 	}
 
@@ -56,7 +57,7 @@ func TestDeadlineHeap_WatchThenUpdateAndGet(t *testing.T) {
 	var batch []string
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(2 * wait):
+	case <-time.After(testutil.Timeout(2 * wait)):
 		t.Fatal("timeout")
 	}
 
@@ -84,7 +85,7 @@ func TestDeadlineHeap_MultiwatchAndDelete(t *testing.T) {
 	var batch []string
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(2 * wait):
+	case <-time.After(testutil.Timeout(2 * wait)):
 		t.Fatal("timeout")
 	}
 
@@ -120,7 +121,7 @@ func TestDeadlineHeap_WatchCoalesce(t *testing.T) {
 	var batch []string
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(1 * time.Second):
+	case <-time.After(testutil.Timeout(time.Second)):
 		t.Fatal("timeout")
 	}
 
@@ -132,7 +133,7 @@ func TestDeadlineHeap_WatchCoalesce(t *testing.T) {
 
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(2 * time.Second):
+	case <-time.After(testutil.Timeout(2 * time.Second)):
 		t.Fatal("timeout")
 	}
 
@@ -144,7 +145,7 @@ func TestDeadlineHeap_WatchCoalesce(t *testing.T) {
 	select {
 	case <-h.NextBatch():
 		t.Fatal("unexpected batch")
-	case <-time.After(100 * time.Millisecond):
+	case <-time.After(testutil.Timeout(100 * time.Millisecond)):
 	}
 }
 
@@ -160,7 +161,7 @@ func TestDeadlineHeap_MultipleForce(t *testing.T) {
 	var batch []string
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(10 * time.Millisecond):
+	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
 		t.Fatal("timeout")
 	}
 
@@ -171,7 +172,7 @@ func TestDeadlineHeap_MultipleForce(t *testing.T) {
 	h.Watch(nodeID, deadline)
 	select {
 	case batch = <-h.NextBatch():
-	case <-time.After(10 * time.Millisecond):
+	case <-time.After(testutil.Timeout(10 * time.Millisecond)):
 		t.Fatal("timeout")
 	}
 
diff --git a/testutil/wait.go b/testutil/wait.go
index f3f993b87..4fd488a52 100644
--- a/testutil/wait.go
+++ b/testutil/wait.go
@@ -62,6 +62,11 @@ func TestMultiplier() int64 {
 	return 1
 }
 
+// Timeout takes the desired timeout and increases it if running in Travis
+func Timeout(original time.Duration) time.Duration {
+	return original * time.Duration(TestMultiplier())
+}
+
 func IsTravis() bool {
 	_, ok := os.LookupEnv(TravisRunEnv)
 	return ok