From c8bc839c816f609d71263a235a8aae0638ff5cd6 Mon Sep 17 00:00:00 2001
From: Alejandro Cabeza Romero <alex93cabeza@gmail.com>
Date: Wed, 7 Aug 2024 21:06:13 +0200
Subject: [PATCH] Refactor discovery tests for reusability.

---
 tests/testdiscovery.nim | 47 ++++++++++++++++++++++++++++-------------
 1 file changed, 32 insertions(+), 15 deletions(-)

diff --git a/tests/testdiscovery.nim b/tests/testdiscovery.nim
index 64ea5304b..c9a45903a 100644
--- a/tests/testdiscovery.nim
+++ b/tests/testdiscovery.nim
@@ -18,7 +18,7 @@ import
     discovery/discoverymngr,
     discovery/rendezvousinterface,
   ]
-import ./helpers
+import ./helpers, ./asyncunit
 
 proc createSwitch(rdv: RendezVous = RendezVous.new()): Switch =
   SwitchBuilder
@@ -32,17 +32,30 @@ proc createSwitch(rdv: RendezVous = RendezVous.new()): Switch =
   .build()
 
 suite "Discovery":
-  teardown:
-    checkTrackers()
-  asyncTest "RendezVous test":
-    let
-      rdvA = RendezVous.new()
-      rdvB = RendezVous.new()
-      clientA = createSwitch(rdvA)
-      clientB = createSwitch(rdvB)
-      remoteNode = createSwitch()
-      dmA = DiscoveryManager()
-      dmB = DiscoveryManager()
+  const
+    namespace = "foo"
+    rdvNamespace = RdvNamespace(namespace)
+
+  var
+    rdvA {.threadvar.}: RendezVous
+    rdvB {.threadvar.}: RendezVous
+    rdvRemote {.threadvar.}: RendezVous
+    clientA {.threadvar.}: Switch
+    clientB {.threadvar.}: Switch
+    remoteNode {.threadvar.}: Switch
+    dmA {.threadvar.}: DiscoveryManager
+    dmB {.threadvar.}: DiscoveryManager
+
+  asyncSetup:
+    rdvA = RendezVous.new()
+    rdvB = RendezVous.new()
+    rdvRemote = RendezVous.new()
+    clientA = createSwitch(rdvA)
+    clientB = createSwitch(rdvB)
+    remoteNode = createSwitch(rdvRemote)
+    dmA = DiscoveryManager()
+    dmB = DiscoveryManager()
+
     dmA.add(RendezVousInterface.new(rdvA, ttr = 500.milliseconds))
     dmB.add(RendezVousInterface.new(rdvB))
     await allFutures(clientA.start(), clientB.start(), remoteNode.start())
@@ -50,14 +63,18 @@ suite "Discovery":
     await clientB.connect(remoteNode.peerInfo.peerId, remoteNode.peerInfo.addrs)
     await clientA.connect(remoteNode.peerInfo.peerId, remoteNode.peerInfo.addrs)
 
-    dmB.advertise(RdvNamespace("foo"))
+  asyncTeardown:
+    await allFutures(clientA.stop(), clientB.stop(), remoteNode.stop())
+    checkTrackers()
+
+  asyncTest "RendezVous test":
+    dmB.advertise(rdvNamespace)
 
     let
-      query = dmA.request(RdvNamespace("foo"))
+      query = dmA.request(rdvNamespace)
       res = await query.getPeer()
     check:
       res{PeerId}.get() == clientB.peerInfo.peerId
       res[PeerId] == clientB.peerInfo.peerId
       res.getAll(PeerId) == @[clientB.peerInfo.peerId]
       toHashSet(res.getAll(MultiAddress)) == toHashSet(clientB.peerInfo.addrs)
-    await allFutures(clientA.stop(), clientB.stop(), remoteNode.stop())