From 8fe44e35c219d4296619b19412eba7ff1a45296f Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Mon, 13 Dec 2021 11:25:33 +0100
Subject: [PATCH] Fix flaky tests (#669)

* Fix push identify flaky test
* Fix mplex flaky tests
* Fix switch flaky test
---
 tests/testidentify.nim |  9 +++++++++
 tests/testmplex.nim    | 10 +++++++++-
 tests/testswitch.nim   |  2 +-
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/tests/testidentify.nim b/tests/testidentify.nim
index 481d5e6..66aca25 100644
--- a/tests/testidentify.nim
+++ b/tests/testidentify.nim
@@ -176,6 +176,9 @@ suite "Identify":
 
       await identifyPush2.push(switch2.peerInfo, conn)
 
+      check await checkExpiring(switch1.peerStore.protoBook.get(switch2.peerInfo.peerId) == switch2.peerInfo.protocols.toHashSet())
+      check await checkExpiring(switch1.peerStore.addressBook.get(switch2.peerInfo.peerId) == switch2.peerInfo.addrs.toHashSet())
+
       await closeAll()
 
       # Wait the very end to be sure that the push has been processed
@@ -197,6 +200,12 @@ suite "Identify":
 
       await identifyPush2.push(switch2.peerInfo, conn)
 
+      # We have no way to know when the message will is received
+      # because it will fail validation inside push identify itself
+      #
+      # So no choice but to sleep
+      await sleepAsync(10.milliseconds)
+
       await closeAll()
 
       # Wait the very end to be sure that the push has been processed
diff --git a/tests/testmplex.nim b/tests/testmplex.nim
index 6141012..774a382 100644
--- a/tests/testmplex.nim
+++ b/tests/testmplex.nim
@@ -816,6 +816,8 @@ suite "Mplex":
       for i in 0..9:
         dialStreams.add((await mplexDial.newStream()))
 
+      check await checkExpiring(listenStreams.len == 10 and dialStreams.len == 10)
+
       await mplexListen.close()
       await allFuturesThrowing(
           (dialStreams & listenStreams)
@@ -860,6 +862,8 @@ suite "Mplex":
       for i in 0..9:
         dialStreams.add((await mplexDial.newStream()))
 
+      check await checkExpiring(listenStreams.len == 10 and dialStreams.len == 10)
+
       mplexHandle.cancel()
       await allFuturesThrowing(
           (dialStreams & listenStreams)
@@ -901,11 +905,14 @@ suite "Mplex":
       for i in 0..9:
         dialStreams.add((await mplexDial.newStream()))
 
+      check await checkExpiring(listenStreams.len == 10 and dialStreams.len == 10)
+
       await conn.close()
       await allFuturesThrowing(
           (dialStreams & listenStreams)
           .mapIt( it.join() ))
 
+
       checkTracker(LPChannelTrackerName)
 
       await conn.closeWithEOF()
@@ -944,7 +951,8 @@ suite "Mplex":
       for i in 0..9:
         dialStreams.add((await mplexDial.newStream()))
 
-      await sleepAsync(100.millis)
+      check await checkExpiring(listenStreams.len == 10 and dialStreams.len == 10)
+
       await listenConn.closeWithEOF()
       await allFuturesThrowing(
           (dialStreams & listenStreams)
diff --git a/tests/testswitch.nim b/tests/testswitch.nim
index c88e690..70c470b 100644
--- a/tests/testswitch.nim
+++ b/tests/testswitch.nim
@@ -634,7 +634,7 @@ suite "Switch":
     await switch.start()
 
     var peerId = PeerID.init(PrivateKey.random(ECDSA, rng[]).get()).get()
-    expect LPStreamClosedError:
+    expect LPStreamClosedError, LPStreamEOFError:
       await switch.connect(peerId, transport.addrs)
 
     await handlerWait