swarm/network: fix TestHiveStatePersistance (#1320)

swarm/network/hive_test.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/p2p"
 	p2ptest "github.com/ethereum/go-ethereum/p2p/testing"
 	"github.com/ethereum/go-ethereum/swarm/state"
 )
@@ -109,67 +110,68 @@ func TestRegisterAndConnect(t *testing.T) {
 // Actual connectivity is not in scope for this test, as the peers loaded from state are not known to
 // the simulation; the test only verifies that the peers are known to the node
 func TestHiveStatePersistance(t *testing.T) {
-
 	dir, err := ioutil.TempDir("", "hive_test_store")
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer os.RemoveAll(dir)
 
-	store, err := state.NewDBStore(dir) //start the hive with an empty dbstore
+	const peersCount = 5
-	if err != nil {
+
-		t.Fatal(err)
+	startHive := func(t *testing.T, dir string) (h *Hive) {
+		store, err := state.NewDBStore(dir)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		params := NewHiveParams()
+		params.Discovery = false
+
+		prvkey, err := crypto.GenerateKey()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		h = NewHive(params, NewKademlia(PrivateKeyToBzzKey(prvkey), NewKadParams()), store)
+		s := p2ptest.NewProtocolTester(prvkey, 0, func(p *p2p.Peer, rw p2p.MsgReadWriter) error { return nil })
+
+		if err := h.Start(s.Server); err != nil {
+			t.Fatal(err)
+		}
+		return h
 	}
 
-	params := NewHiveParams()
+	h1 := startHive(t, dir)
-	params.Discovery = false
-
-	s, pp, err := newHiveTester(t, params, 5, store)
-	if err != nil {
-		t.Fatal(err)
-	}
 	peers := make(map[string]bool)
-	for _, node := range s.Nodes {
+	for i := 0; i < peersCount; i++ {
-		raddr := NewAddr(node)
+		raddr := RandomAddr()
-		pp.Register(raddr)
+		h1.Register(raddr)
-		log.Warn("add", "addr", raddr.String())
+		log.Trace("add", "addr", raddr.String())
 		peers[raddr.String()] = true
 	}
-
+	if err = h1.Stop(); err != nil {
-	// start and stop the hive
-	// the known peers should be saved upon stopping
-	err = pp.Start(s.Server)
-	if err != nil {
-		t.Fatal(err)
-	}
-	pp.Stop()
-	store.Close()
-
-	// start the hive with an empty dbstore
-	persistedStore, err := state.NewDBStore(dir)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	s1, pp, err := newHiveTester(t, params, 0, persistedStore)
-	if err != nil {
 		t.Fatal(err)
 	}
 
 	// start the hive and check that we know of all expected peers
-	pp.Start(s1.Server)
+	h2 := startHive(t, dir)
+	defer func() {
+		if err = h2.Stop(); err != nil {
+			t.Fatal(err)
+		}
+	}()
+
 	i := 0
-	pp.Kademlia.EachAddr(nil, 256, func(addr *BzzAddr, po int) bool {
+	h2.Kademlia.EachAddr(nil, 256, func(addr *BzzAddr, po int) bool {
-		log.Warn("check", "addr", addr.String())
+		log.Trace("check", "addr", addr.String())
 		delete(peers, addr.String())
 		i++
 		return true
 	})
-	if i != 5 {
+	if i != peersCount {
-		t.Fatalf("invalid number of entries: got %v, want %v", i, 5)
+		t.Fatalf("invalid number of entries: got %v, want %v", i, peersCount)
 	}
 	if len(peers) != 0 {
 		t.Fatalf("%d peers left over: %v", len(peers), peers)
 	}
-
 }