swarm/network/stream: refactored TestGetSubscriptionsRPC

swarm/network/stream/streamer_test.go

@@ -23,15 +23,21 @@ import (
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	p2ptest "github.com/ethereum/go-ethereum/p2p/testing"
 	"github.com/ethereum/go-ethereum/swarm/network"
+	"github.com/ethereum/go-ethereum/swarm/network/simulation"
+	"github.com/ethereum/go-ethereum/swarm/state"
+	"github.com/ethereum/go-ethereum/swarm/storage"
 	"golang.org/x/crypto/sha3"
 )
 
@@ -1120,19 +1126,31 @@ func TestRequestPeerSubscriptions(t *testing.T) {
 TestGetSubscriptionsRPC sets up a simulation network of 16 nodes,
 starts the simulation, waits for SyncUpdateDelay in order to kick off
 stream registration, then tests that there are subscriptions.
-If provided with the `-printstats = true` option, it will print
-the information of who is subscribed to who to STDOUT
 */
 func TestGetSubscriptionsRPC(t *testing.T) {
 	// arbitrarily set to 16
 	nodeCount := 16
 	// set the syncUpdateDelay for sync registrations to start
 	syncUpdateDelay := 500 * time.Millisecond
-	//we will later need the kad table for each node
-	bucketKeyKad := simulation.BucketKey("kademlia")
 	// holds the msg code for SubscribeMsg
 	var subscribeMsgCode uint64
 	var ok bool
+	var expectedMsgCount = 0
+
+	// this channel signalizes that the expected amount of subscriptiosn is done
+	allSubscriptionsDone := make(chan struct{})
+	lock := sync.RWMutex{}
+	// after the test, we need to reset the subscriptionFunc to the default
+	defer func() { subscriptionFunc = doRequestSubscription }()
+
+	// we use this subscriptionFunc for this test: just increases count and calls the actual subscription
+	subscriptionFunc = func(r *Registry, p *network.Peer, bin uint8, subs map[enode.ID]map[Stream]struct{}) bool {
+		lock.Lock()
+		expectedMsgCount++
+		lock.Unlock()
+		doRequestSubscription(r, p, bin, subs)
+		return true
+	}
 	// create a standard sim
 	sim := simulation.New(map[string]simulation.ServiceFunc{
 		"streamer": func(ctx *adapters.ServiceContext, bucket *sync.Map) (s node.Service, cleanup func(), err error) {
@@ -1148,8 +1166,6 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 				return nil, nil, err
 			}
 			kad := network.NewKademlia(addr.Over(), network.NewKadParams())
-			//store the kad table
-			bucket.Store(bucketKeyKad, kad)
 			delivery := NewDelivery(kad, netStore)
 			netStore.NewNetFetcherFunc = network.NewFetcherFactory(dummyRequestFromPeers, true).New
 			// configure so that sync registrations actually happen
@@ -1164,7 +1180,6 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 				t.Fatal("Message code for SubscribeMsg not found")
 			}
 
-			bucket.Store(bucketKeyRegistry, r)
 			cleanup = func() {
 				os.RemoveAll(datadir)
 				netStore.Close()
@@ -1186,10 +1201,6 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	//run the simulation
-	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
-		log.Info("Simulation running")
-		nodes := sim.Net.Nodes
 	// setup the filter for SubscribeMsg
 	msgs := sim.PeerEvents(
 		context.Background(),
@@ -1197,47 +1208,39 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 		simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode),
 	)
 
-		//setup the vars we need
+	// strategy: listen to all SubscribeMsg events; after every event we wait
-		msgCount := 0
+	// if after 1 second no more messages are being received, we assume the
-		expectedMsgCount := 0
+	// subscription phase has terminated!
-		allSubscriptionsDone := make(chan struct{})
 
-		//in a loop, catch all SubscribeMsg and just add up
+	// the loop in this go routine will either wait for new message events
+	// or times out after 1 second, which signals that we are not receiving
+	// any new subscriptions any more
 	go func() {
-			for m := range msgs {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case m := <-msgs: // just reset the loop
 				if m.Error != nil {
 					log.Error("stream message", "err", m.Error)
 					continue
 				}
 				log.Trace("stream message", "node", m.NodeID, "peer", m.PeerID)
-				//add one
+			case <-time.After(time.Second):
-				msgCount += 1
+				// one second passed, don't assume more subscriptions
-				if msgCount == expectedMsgCount {
-					//the expected amount is reached
 				allSubscriptionsDone <- struct{}{}
+				log.Info("All subscriptions received")
 				return
+
 			}
 		}
 	}()
 
-		//first iterate all nodes to get the expected number of subscriptions from the kad table
+	//run the simulation
-		for _, node := range nodes {
+	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
-			item, ok := sim.NodeItem(node.ID(), bucketKeyKad)
+		log.Info("Simulation running")
-			if !ok {
+		nodes := sim.Net.Nodes
-				return fmt.Errorf("No kademlia")
+
-			}
-			kad := item.(*network.Kademlia)
-			//define the function which should run for each connection - just count subscriptions
-			//this is not actually subscribing but iterating the same way as the subscriptions do,
-			//as we need just the number
-			eachBinFunc := func(p *network.Peer, bin int) bool {
-				expectedMsgCount += 1
-				return true
-			}
-			//call the actual kademlia for the count
-			kad.EachBin(kad.BaseAddr(), pot.DefaultPof(kad.MaxProxDisplay), 0, eachBinFunc)
-		}
-		log.Debug("Expected message count: ", "expectedMsgCount", expectedMsgCount)
 		//wait until all subscriptions are done
 		select {
 		case <-allSubscriptionsDone:
@@ -1245,8 +1248,9 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 			t.Fatal("Context timed out")
 		}
 
-		log.Info("All subscriptions received")
+		log.Debug("Expected message count: ", "expectedMsgCount", expectedMsgCount)
 		//now iterate again, this time we call each node via RPC to get its subscriptions
+		realCount := 0
 		for _, node := range nodes {
 			//create rpc client
 			client, err := node.Client()
@@ -1254,11 +1258,6 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 				t.Fatalf("create node 1 rpc client fail: %v", err)
 			}
 
-			item, ok := sim.NodeItem(node.ID(), bucketKeyRegistry)
-			if !ok {
-				return fmt.Errorf("No registry")
-			}
-			registry := item.(*Registry)
 			//ask it for subscriptions
 			pstreams := make(map[string][]string)
 			err = client.Call(&pstreams, "stream_getPeerSubscriptions")
@@ -1266,17 +1265,23 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 				t.Fatal(err)
 			}
 			//length of the subscriptions can not be smaller than number of peers
-			if len(pstreams) < len(registry.peers) {
-				t.Fatal("Subscription count is smaller than expected")
-			}
 			log.Debug(fmt.Sprintf("node %s subscriptions:", node.String()))
 			for p, ps := range pstreams {
 				log.Debug(fmt.Sprintf("...with node %s: ", p))
 				for _, s := range ps {
 					log.Debug(fmt.Sprintf("......%s", s))
+					// each node also has subscriptions to RETRIEVE_REQUEST streams,
+					// we need to ignore those, we are only counting SYNC streams
+					if !strings.HasPrefix(s, "RETRIEVE_REQUEST") {
+						realCount++
 					}
 				}
 			}
+		}
+		// every node is mutually subscribed to each other, so the actual count is half of it
+		if realCount/2 != expectedMsgCount {
+			return errors.New(fmt.Sprintf("Real subscriptions and expected amount don't match; real: %d, expected: %d", realCount/2, expectedMsgCount))
+		}
 		return nil
 	})
 	if result.Error != nil {