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,20 +1126,32 @@ 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
+	// arbitrarily set to 16
 	nodeCount := 16
-	//set the syncUpdateDelay for sync registrations to start
+	// 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
+	// holds the msg code for SubscribeMsg
 	var subscribeMsgCode uint64
 	var ok bool
-	//create a standard sim
+	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) {
 			n := ctx.Config.Node()
@@ -1148,23 +1166,20 @@ 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
+			// configure so that sync registrations actually happen
 			r := NewRegistry(addr.ID(), delivery, netStore, state.NewInmemoryStore(), &RegistryOptions{
 				Retrieval:       RetrievalEnabled,
 				Syncing:         SyncingAutoSubscribe, //enable sync registrations
 				SyncUpdateDelay: syncUpdateDelay,
 			}, nil)
-			//get the SubscribeMsg code
+			// get the SubscribeMsg code
 			subscribeMsgCode, ok = r.GetSpec().GetCode(SubscribeMsg{})
 			if !ok {
 				t.Fatal("Message code for SubscribeMsg not found")
 			}
 
-			bucket.Store(bucketKeyRegistry, r)
 			cleanup = func() {
 				os.RemoveAll(datadir)
 				netStore.Close()
@@ -1180,64 +1195,52 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 	ctx, cancelSimRun := context.WithTimeout(context.Background(), 1*time.Minute)
 	defer cancelSimRun()
 
-	//upload a snapshot
+	// upload a snapshot
 	err := sim.UploadSnapshot(fmt.Sprintf("testing/snapshot_%d.json", nodeCount))
 	if err != nil {
 		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(),
-			sim.NodeIDs(),
-			simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode),
-		)
+	// setup the filter for SubscribeMsg
+	msgs := sim.PeerEvents(
+		context.Background(),
+		sim.NodeIDs(),
+		simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode),
+	)
 
-		//setup the vars we need
-		msgCount := 0
-		expectedMsgCount := 0
-		allSubscriptionsDone := make(chan struct{})
+	// strategy: listen to all SubscribeMsg events; after every event we wait
+	// if after 1 second no more messages are being received, we assume the
+	// subscription phase has terminated!
 
-		//in a loop, catch all SubscribeMsg and just add up
-		go func() {
-			for m := range msgs {
+	// 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 {
+			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
-				msgCount += 1
-				if msgCount == expectedMsgCount {
-					//the expected amount is reached
-					allSubscriptionsDone <- struct{}{}
-					return
-				}
-			}
-		}()
+			case <-time.After(time.Second):
+				// one second passed, don't assume more subscriptions
+				allSubscriptionsDone <- struct{}{}
+				log.Info("All subscriptions received")
+				return
 
-		//first iterate all nodes to get the expected number of subscriptions from the kad table
-		for _, node := range nodes {
-			item, ok := sim.NodeItem(node.ID(), bucketKeyKad)
-			if !ok {
-				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)
+	}()
+
+	//run the simulation
+	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
+		log.Info("Simulation running")
+		nodes := sim.Net.Nodes
+
 		//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 {