From 2eaa566371675fc38d5a8a0939535458e08bc3d4 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Wed, 22 Apr 2020 16:58:07 +0300
Subject: [PATCH] add direct test for gossipsub router's piggybackControl

---
 gossipsub_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/gossipsub_test.go b/gossipsub_test.go
index 97a0c0f..88dda36 100644
--- a/gossipsub_test.go
+++ b/gossipsub_test.go
@@ -8,9 +8,14 @@ import (
 	"testing"
 	"time"
 
+	pb "github.com/libp2p/go-libp2p-pubsub/pb"
+
 	"github.com/libp2p/go-libp2p-core/host"
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
+
+	bhost "github.com/libp2p/go-libp2p-blankhost"
+	swarmt "github.com/libp2p/go-libp2p-swarm/testing"
 )
 
 func getGossipsub(ctx context.Context, h host.Host, opts ...Option) *PubSub {
@@ -1214,3 +1219,53 @@ func TestGossipsubNegativeScore(t *testing.T) {
 		}
 	}
 }
+
+func TestGossipsubPiggybackControl(t *testing.T) {
+	// this is a direct test of the piggybackControl function as we can't reliably
+	// trigger it on travis
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	h := bhost.NewBlankHost(swarmt.GenSwarm(t, ctx))
+	ps := getGossipsub(ctx, h)
+
+	blah := peer.ID("bogotr0n")
+
+	res := make(chan *RPC, 1)
+	ps.eval <- func() {
+		gs := ps.rt.(*GossipSubRouter)
+		test1 := "test1"
+		test2 := "test2"
+		test3 := "test3"
+		gs.mesh[test1] = make(map[peer.ID]struct{})
+		gs.mesh[test2] = make(map[peer.ID]struct{})
+		gs.mesh[test1][blah] = struct{}{}
+
+		rpc := &RPC{RPC: pb.RPC{}}
+		gs.piggybackControl(blah, rpc, &pb.ControlMessage{
+			Graft: []*pb.ControlGraft{&pb.ControlGraft{TopicID: &test1}, &pb.ControlGraft{TopicID: &test2}, &pb.ControlGraft{TopicID: &test3}},
+			Prune: []*pb.ControlPrune{&pb.ControlPrune{TopicID: &test1}, &pb.ControlPrune{TopicID: &test2}, &pb.ControlPrune{TopicID: &test3}},
+		})
+		res <- rpc
+	}
+
+	rpc := <-res
+	if rpc.Control == nil {
+		t.Fatal("expected non-nil control message")
+	}
+	if len(rpc.Control.Graft) != 1 {
+		t.Fatal("expected 1 GRAFT")
+	}
+	if rpc.Control.Graft[0].GetTopicID() != "test1" {
+		t.Fatal("expected test1 as graft topic ID")
+	}
+	if len(rpc.Control.Prune) != 2 {
+		t.Fatal("expected 2 PRUNEs")
+	}
+	if rpc.Control.Prune[0].GetTopicID() != "test2" {
+		t.Fatal("expected test2 as prune topic ID")
+	}
+	if rpc.Control.Prune[1].GetTopicID() != "test3" {
+		t.Fatal("expected test3 as prune topic ID")
+	}
+}