From 640ebdaef6e7787e239bee9ad5d0aff93ad19e2a Mon Sep 17 00:00:00 2001
From: Alex Dadgar <alex.dadgar@gmail.com>
Date: Mon, 19 Mar 2018 15:19:57 -0700
Subject: [PATCH] fix race in drain integration tests

---
 nomad/drainer_int_test.go | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/nomad/drainer_int_test.go b/nomad/drainer_int_test.go
index 8e03a2ef5..f71363a0d 100644
--- a/nomad/drainer_int_test.go
+++ b/nomad/drainer_int_test.go
@@ -60,7 +60,13 @@ func allocPromoter(t *testing.T, ctx context.Context,
 			WriteRequest: structs.WriteRequest{Region: "global"},
 		}
 		var resp structs.NodeAllocsResponse
-		require.Nil(t, msgpackrpc.CallWithCodec(codec, "Node.UpdateAlloc", req, &resp))
+		if err := msgpackrpc.CallWithCodec(codec, "Node.UpdateAlloc", req, &resp); err != nil {
+			if ctx.Err() == context.Canceled {
+				return
+			} else {
+				require.Nil(t, err)
+			}
+		}
 	}
 }