From bb6a3c2ce1716175a7e22918898e0a74421140eb Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Mon, 17 Feb 2020 16:36:04 +0100
Subject: [PATCH] Add basic node discovery test

---
 eth.nimble                       |  1 +
 eth/p2p/discoveryv5/protocol.nim |  2 +-
 tests/p2p/test_discoveryv5.nim   | 52 ++++++++++++++++++++++++++++++++
 3 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 tests/p2p/test_discoveryv5.nim

diff --git a/eth.nimble b/eth.nimble
index c6ab841..a09a67c 100644
--- a/eth.nimble
+++ b/eth.nimble
@@ -56,6 +56,7 @@ proc runP2pTests() =
       "test_waku_mail",
       "test_protocol_handlers",
       "test_enr",
+      "test_discoveryv5"
     ]:
     runTest("tests/p2p/" & filename)
 
diff --git a/eth/p2p/discoveryv5/protocol.nim b/eth/p2p/discoveryv5/protocol.nim
index 60e5b6c..e428a55 100644
--- a/eth/p2p/discoveryv5/protocol.nim
+++ b/eth/p2p/discoveryv5/protocol.nim
@@ -8,7 +8,7 @@ import nimcrypto except toHex
 type
   Protocol* = ref object
     transp: DatagramTransport
-    localNode: Node
+    localNode*: Node
     privateKey: PrivateKey
     whoareyouMagic: array[32, byte]
     idHash: array[32, byte]
diff --git a/tests/p2p/test_discoveryv5.nim b/tests/p2p/test_discoveryv5.nim
new file mode 100644
index 0000000..4e2776c
--- /dev/null
+++ b/tests/p2p/test_discoveryv5.nim
@@ -0,0 +1,52 @@
+import
+  unittest, chronos, sequtils, chronicles,
+  eth/keys, eth/p2p/enode, eth/trie/db,
+  eth/p2p/discoveryv5/[discovery_db, enr, node, types],
+  eth/p2p/discoveryv5/protocol as discv5_protocol,
+  ./p2p_test_helper
+
+proc startDiscoveryv5Node*(privKey: PrivateKey, address: Address,
+                        bootnodes: seq[Record]): discv5_protocol.Protocol =
+  var db = DiscoveryDB.init(newMemoryDB())
+  result = newProtocol(privKey, db, address.udpPort)
+
+  for node in bootnodes:
+    result.addNode(node)
+
+  result.open()
+  result.start()
+
+proc nodeIdInNodes(id: NodeId, nodes: openarray[Node]): bool =
+  for n in nodes:
+    if id == n.id: return true
+
+suite "Discovery v5 Tests":
+  asyncTest "Discover nodes":
+    let
+      bootNodeKey = initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a617")
+      bootNodeAddr = localAddress(20301)
+      bootNode = startDiscoveryv5Node(bootNodeKey, bootNodeAddr, @[])
+      bootNodeRecord = initRecord(1, bootNodeKey,
+        {"udp": bootNodeAddr.udpPort.uint16, "ip": [byte 127, 0, 0, 1]})
+
+    let nodeKeys = [
+        initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a618"),
+        initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a619"),
+        initPrivateKey("a2b50376a79b1a8c8a3296485572bdfbf54708bb46d3c25d73d2723aaaf6a620")
+      ]
+    var nodeAddrs = newSeqOfCap[Address](nodeKeys.len)
+    for i in 0 ..< nodeKeys.len: nodeAddrs.add(localAddress(20302 + i))
+
+    var nodes = zip(nodeKeys, nodeAddrs).mapIt(
+      startDiscoveryv5Node(it.a, it.b, @[bootNodeRecord]))
+    nodes.add(bootNode)
+
+    for node in nodes:
+      let discovered = await node.lookupRandom()
+      debug "Lookup from random id", node=node.localNode, discovered
+
+    # Check for each node if the other nodes shows up in the routing table
+    for i in nodes:
+      for j in nodes:
+        if j != i:
+          check(nodeIdInNodes(i.localNode.id, j.randomNodes(nodes.len - 1)))