From c9a9b02e489bdc5fe8b7fa8f2eb4cfa16491838b Mon Sep 17 00:00:00 2001
From: Richard Ramos <info@richardramos.me>
Date: Thu, 18 Nov 2021 08:26:36 -0400
Subject: [PATCH] test: stop / start discovery (#153)

---
 waku/v2/discv5/discover_test.go | 80 +++++++++++++++++++++++++++++----
 1 file changed, 72 insertions(+), 8 deletions(-)

diff --git a/waku/v2/discv5/discover_test.go b/waku/v2/discv5/discover_test.go
index f4d8bfa1..f9d6d448 100644
--- a/waku/v2/discv5/discover_test.go
+++ b/waku/v2/discv5/discover_test.go
@@ -72,24 +72,88 @@ func TestDiscV5(t *testing.T) {
 	err = d3.Start()
 	require.NoError(t, err)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancel()
 
-	peerChan, err := d1.FindPeers(ctx, "", discovery.Limit(2))
+	peerChan, err := d3.FindPeers(ctx, "", discovery.Limit(2))
 	require.NoError(t, err)
 
+	foundHost1 := false
 	foundHost2 := false
-	foundHost3 := false
 	for p := range peerChan {
-		if p.Addrs[0].String() == host2.Addrs()[0].String() {
-			foundHost2 = true
+		if p.Addrs[0].String() == host1.Addrs()[0].String() {
+			foundHost1 = true
 		}
 
-		if p.Addrs[0].String() == host3.Addrs()[0].String() {
-			foundHost3 = true
+		if p.Addrs[0].String() == host2.Addrs()[0].String() {
+			foundHost2 = true
 
 		}
 	}
 
-	require.True(t, foundHost2 && foundHost3)
+	require.True(t, foundHost1 && foundHost2)
+
+	// Should return nodes from the cache
+
+	d3.Stop()
+
+	foundHost1 = false
+	foundHost2 = false
+
+	ctx1, cancel1 := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel1()
+
+	peerChan, err = d3.FindPeers(ctx1, "", discovery.Limit(2))
+	require.NoError(t, err)
+	for p := range peerChan {
+		if p.Addrs[0].String() == host1.Addrs()[0].String() {
+			foundHost1 = true
+		}
+
+		if p.Addrs[0].String() == host2.Addrs()[0].String() {
+			foundHost2 = true
+		}
+	}
+
+	require.True(t, foundHost1 && foundHost2)
+
+	// Simulate empty cache
+
+	for i := range d3.peerCache.recs {
+		delete(d3.peerCache.recs, i)
+	}
+
+	ctx2, cancel2 := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel2()
+
+	peerChan, err = d3.FindPeers(ctx2, "", discovery.Limit(2))
+	require.NoError(t, err)
+	for range peerChan {
+		require.Fail(t, "Should not have peers")
+	}
+
+	// Restart peer search
+	err = d3.Start()
+	require.NoError(t, err)
+
+	foundHost1 = false
+	foundHost2 = false
+
+	ctx3, cancel3 := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel3()
+
+	peerChan, err = d3.FindPeers(ctx3, "", discovery.Limit(2))
+	require.NoError(t, err)
+	for p := range peerChan {
+		if p.Addrs[0].String() == host1.Addrs()[0].String() {
+			foundHost1 = true
+		}
+
+		if p.Addrs[0].String() == host2.Addrs()[0].String() {
+			foundHost2 = true
+		}
+	}
+
+	require.True(t, foundHost1 && foundHost2)
+
 }