From 4037390c421ed18132514beee91c5b9f0d10463d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ra=C3=BAl=20Kripalani?= <raulk@apache.org>
Date: Tue, 11 Sep 2018 18:10:27 +0100
Subject: [PATCH] increase test coverage.

---
 pstoreds/addr_book.go   |  1 -
 pstoreds/ds_test.go     | 21 +++++++++++++---
 test/addr_book_suite.go | 55 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/pstoreds/addr_book.go b/pstoreds/addr_book.go
index 53c3006..9bb5a45 100644
--- a/pstoreds/addr_book.go
+++ b/pstoreds/addr_book.go
@@ -321,7 +321,6 @@ func (mgr *dsAddrBook) ClearAddrs(p peer.ID) {
 	if e, ok := mgr.cache.Peek(p.Pretty()); ok {
 		mgr.cache.Remove(p.Pretty())
 		keys, _, _ := keysAndAddrs(p, e.([]ma.Multiaddr))
-
 		deleteFn = func() error {
 			return mgr.dbDelete(keys)
 		}
diff --git a/pstoreds/ds_test.go b/pstoreds/ds_test.go
index f1972dc..4a8d04c 100644
--- a/pstoreds/ds_test.go
+++ b/pstoreds/ds_test.go
@@ -122,10 +122,25 @@ func TestBadgerDsPeerstore(t *testing.T) {
 }
 
 func TestBadgerDsAddrBook(t *testing.T) {
-	opts := DefaultOpts()
-	opts.TTLInterval = 100 * time.Microsecond
+	t.Run("Cacheful", func(t *testing.T) {
+		t.Parallel()
 
-	pt.TestAddrBook(t, addressBookFactory(t, opts))
+		opts := DefaultOpts()
+		opts.TTLInterval = 100 * time.Microsecond
+		opts.CacheSize = 1024
+
+		pt.TestAddrBook(t, addressBookFactory(t, opts))
+	})
+
+	t.Run("Cacheless", func(t *testing.T) {
+		t.Parallel()
+
+		opts := DefaultOpts()
+		opts.TTLInterval = 100 * time.Microsecond
+		opts.CacheSize = 0
+
+		pt.TestAddrBook(t, addressBookFactory(t, opts))
+	})
 }
 
 func BenchmarkBadgerDsPeerstore(b *testing.B) {
diff --git a/test/addr_book_suite.go b/test/addr_book_suite.go
index 441e2fa..dbff04a 100644
--- a/test/addr_book_suite.go
+++ b/test/addr_book_suite.go
@@ -18,6 +18,8 @@ var addressBookSuite = map[string]func(book pstore.AddrBook) func(*testing.T){
 	"UpdateTTLs":           testUpdateTTLs,
 	"NilAddrsDontBreak":    testNilAddrsDontBreak,
 	"AddressesExpire":      testAddressesExpire,
+	"ClearWithIter":        testClearWithIterator,
+	"PeersWithAddresses":   testPeersWithAddrs,
 }
 
 type AddrBookFactory func() (pstore.AddrBook, func())
@@ -234,6 +236,59 @@ func testAddressesExpire(m pstore.AddrBook) func(t *testing.T) {
 	}
 }
 
+func testClearWithIterator(m pstore.AddrBook) func(t *testing.T) {
+	return func(t *testing.T) {
+		ids := generatePeerIds(2)
+		addrs := generateAddrs(100)
+
+		// Add the peers with 50 addresses each.
+		m.AddAddrs(ids[0], addrs[:50], pstore.PermanentAddrTTL)
+		m.AddAddrs(ids[1], addrs[50:], pstore.PermanentAddrTTL)
+
+		if all := append(m.Addrs(ids[0]), m.Addrs(ids[1])...); len(all) != 100 {
+			t.Fatal("expected pstore to contain both peers with all their maddrs")
+		}
+
+		// Since we don't fetch these peers, they won't be present in cache.
+
+		m.ClearAddrs(ids[0])
+		if all := append(m.Addrs(ids[0]), m.Addrs(ids[1])...); len(all) != 50 {
+			t.Fatal("expected pstore to contain only addrs of peer 2")
+		}
+
+		m.ClearAddrs(ids[1])
+		if all := append(m.Addrs(ids[0]), m.Addrs(ids[1])...); len(all) != 0 {
+			t.Fatal("expected pstore to contain no addresses")
+		}
+	}
+}
+
+func testPeersWithAddrs(m pstore.AddrBook) func(t *testing.T) {
+	return func(t *testing.T) {
+		// cannot run in parallel as the store is modified.
+		// go runs sequentially in the specified order
+		// see https://blog.golang.org/subtests
+
+		t.Run("empty addrbook", func(t *testing.T) {
+			if peers := m.PeersWithAddrs(); len(peers) != 0 {
+				t.Fatal("expected to find no peers")
+			}
+		})
+
+		t.Run("non-empty addrbook", func(t *testing.T) {
+			ids := generatePeerIds(2)
+			addrs := generateAddrs(10)
+
+			m.AddAddrs(ids[0], addrs[:5], pstore.PermanentAddrTTL)
+			m.AddAddrs(ids[1], addrs[5:], pstore.PermanentAddrTTL)
+
+			if peers := m.PeersWithAddrs(); len(peers) != 2 {
+				t.Fatal("expected to find 2 peers")
+			}
+		})
+	}
+}
+
 func testHas(t *testing.T, exp, act []ma.Multiaddr) {
 	t.Helper()
 	if len(exp) != len(act) {