From 5506c2ae4ac14b0644f1ef89a7cefae35bd016cd Mon Sep 17 00:00:00 2001
From: Sonic <sonic@cobinhood.com>
Date: Mon, 15 Oct 2018 12:53:29 +0800
Subject: dex: fix flaky test

---
 dex/protocol_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

(limited to 'dex')

diff --git a/dex/protocol_test.go b/dex/protocol_test.go
index a26a40feb..4bb3dc9e8 100644
--- a/dex/protocol_test.go
+++ b/dex/protocol_test.go
@@ -406,6 +406,7 @@ func TestSendLatticeBlock(t *testing.T) {
 		},
 	}
 
+	waitForRegister(pm, 1)
 	pm.BroadcastLatticeBlock(&block)
 	msg, err := p.app.ReadMsg()
 	if err != nil {
@@ -552,6 +553,7 @@ func TestSendVote(t *testing.T) {
 		wg.Add(1)
 		go checkvote(p, tt.isReceiver)
 	}
+	waitForRegister(pm, len(testPeers))
 	pm.BroadcastVote(&vote)
 	wg.Wait()
 }
@@ -725,6 +727,7 @@ func TestSendAgreement(t *testing.T) {
 		Votes:     []coreTypes.Vote{vote},
 	}
 
+	waitForRegister(pm, 1)
 	pm.BroadcastAgreementResult(&agreement)
 	msg, err := p.app.ReadMsg()
 	if err != nil {
@@ -785,6 +788,7 @@ func TestSendRandomness(t *testing.T) {
 		Randomness: []byte{7, 7, 7, 7},
 	}
 
+	waitForRegister(pm, 1)
 	pm.BroadcastRandomnessResult(&randomness)
 	msg, err := p.app.ReadMsg()
 	if err != nil {
@@ -802,3 +806,11 @@ func TestSendRandomness(t *testing.T) {
 		t.Errorf("agreement mismatch")
 	}
 }
+
+func waitForRegister(pm *ProtocolManager, num int) {
+	for {
+		if pm.peers.Len() >= num {
+			return
+		}
+	}
+}
-- 
cgit v1.2.3