Add tests

Suggested-by: Leo Chen <leo@harmony.one>

p2p/group_test.go

@@ -0,0 +1,22 @@
+package p2p
+
+import "testing"
+
+func TestGroupID_String(t *testing.T) {
+	tests := []struct {
+		name string
+		id   GroupID
+		want string
+	}{
+		{"empty", GroupID(""), ""},
+		{"ABC", GroupID("ABC"), "414243"},
+		{"binary", GroupID([]byte{1, 2, 3}), "010203"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.id.String(); got != tt.want {
+				t.Errorf("GroupID.String() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

p2p/host/hostv2/hostv2.go

@@ -5,8 +5,6 @@ import (
 	"fmt"
 	"io"
 
-	pubsub "github.com/libp2p/go-libp2p-pubsub"
-
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/harmony-one/harmony/p2p"
 	libp2p "github.com/libp2p/go-libp2p"
@@ -16,6 +14,7 @@ import (
 	peer "github.com/libp2p/go-libp2p-peer"
 	peerstore "github.com/libp2p/go-libp2p-peerstore"
 	protocol "github.com/libp2p/go-libp2p-protocol"
+	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	p2p_config "github.com/libp2p/go-libp2p/config"
 	ma "github.com/multiformats/go-multiaddr"
 )
@@ -27,10 +26,16 @@ const (
 	ProtocolID = "/harmony/0.0.1"
 )
 
+// PubSub captures the pubsub interface we expect from libp2p.
+type PubSub interface {
+	Publish(topic string, data []byte) error
+	Subscribe(topic string, opts ...pubsub.SubOpt) (*pubsub.Subscription, error)
+}
+
 // HostV2 is the version 2 p2p host
 type HostV2 struct {
 	h      p2p_host.Host
-	pubsub *pubsub.PubSub
+	pubsub PubSub
 	self   p2p.Peer
 	priKey p2p_crypto.PrivKey
 }
@@ -47,9 +52,15 @@ func (host *HostV2) SendMessageToGroups(groups []p2p.GroupID, msg []byte) error
 	return error
 }
 
+// Subscription captures the subscription interface
+type Subscription interface {
+	Next(ctx context.Context) (*pubsub.Message, error)
+	Cancel()
+}
+
 // GroupReceiver is a multicast group receiver implementation.
 type GroupReceiver struct {
-	sub *pubsub.Subscription
+	sub Subscription
 }
 
 // Close closes this receiver.

p2p/host/hostv2/hostv2_test.go

@@ -0,0 +1,132 @@
+package hostv2
+
+//go:generate mockgen -source hostv2.go -destination=mock/hostv2_mock.go
+
+import (
+	"context"
+	"errors"
+	"reflect"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/harmony-one/harmony/p2p"
+	mock "github.com/harmony-one/harmony/p2p/host/hostv2/mock"
+	peer "github.com/libp2p/go-libp2p-peer"
+	libp2p_pubsub "github.com/libp2p/go-libp2p-pubsub"
+	libp2p_pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
+)
+
+func TestHostV2_SendMessageToGroups(t *testing.T) {
+	t.Run("Basic", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+		groups := []p2p.GroupID{"ABC", "DEF"}
+		data := []byte{1, 2, 3}
+		pubsub := mock.NewMockPubSub(mc)
+		gomock.InOrder(
+			pubsub.EXPECT().Publish("ABC", data),
+			pubsub.EXPECT().Publish("DEF", data),
+		)
+		host := &HostV2{pubsub: pubsub}
+		if err := host.SendMessageToGroups(groups, data); err != nil {
+			t.Errorf("expected no error; got %v", err)
+		}
+	})
+	t.Run("Error", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+		groups := []p2p.GroupID{"ABC", "DEF"}
+		data := []byte{1, 2, 3}
+		pubsub := mock.NewMockPubSub(mc)
+		gomock.InOrder(
+			pubsub.EXPECT().Publish("ABC", data).Return(errors.New("FIAL")),
+			pubsub.EXPECT().Publish("DEF", data), // Should not early-return
+		)
+		host := &HostV2{pubsub: pubsub}
+		if err := host.SendMessageToGroups(groups, data); err == nil {
+			t.Error("expected an error but got none")
+		}
+	})
+}
+
+func TestGroupReceiver_Close(t *testing.T) {
+	mc := gomock.NewController(t)
+	defer mc.Finish()
+	sub := mock.NewMockSubscription(mc)
+	sub.EXPECT().Cancel()
+	receiver := GroupReceiver{sub: sub}
+	if err := receiver.Close(); err != nil {
+		t.Errorf("expected no error but got %v", err)
+	}
+}
+
+func pubsubMessage(from peer.ID, data []byte) *libp2p_pubsub.Message {
+	m := libp2p_pubsub_pb.Message{From: []byte(from), Data: data}
+	return &libp2p_pubsub.Message{Message: &m}
+}
+
+func TestGroupReceiver_Receive(t *testing.T) {
+	mc := gomock.NewController(t)
+	defer mc.Finish()
+	sub := mock.NewMockSubscription(mc)
+	ctx, _ := context.WithCancel(context.Background())
+	gomock.InOrder(
+		sub.EXPECT().Next(ctx).Return(pubsubMessage("ABC", []byte{1, 2, 3}), nil),
+		sub.EXPECT().Next(ctx).Return(pubsubMessage("DEF", []byte{4, 5, 6}), nil),
+		sub.EXPECT().Next(ctx).Return(nil, errors.New("FIAL")),
+	)
+	receiver := GroupReceiver{sub: sub}
+	verify := func(sender peer.ID, msg []byte, shouldError bool) {
+		gotMsg, gotSender, err := receiver.Receive(ctx)
+		if (err != nil) != shouldError {
+			if shouldError {
+				t.Error("expected an error but got none")
+			} else {
+				t.Errorf("expected no error but got %v", err)
+			}
+		}
+		if gotSender != sender {
+			t.Errorf("expected sender %v but got %v", sender, gotSender)
+		}
+		if !reflect.DeepEqual(gotMsg, msg) {
+			t.Errorf("expected message %v but got %v", msg, gotMsg)
+		}
+	}
+	verify("ABC", []byte{1, 2, 3}, false)
+	verify("DEF", []byte{4, 5, 6}, false)
+	verify("", nil, true)
+}
+
+func TestHostV2_GroupReceiver(t *testing.T) {
+	t.Run("Basic", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+		sub := &libp2p_pubsub.Subscription{}
+		pubsub := mock.NewMockPubSub(mc)
+		pubsub.EXPECT().Subscribe("ABC").Return(sub, nil)
+		host := &HostV2{pubsub: pubsub}
+		gotReceiver, err := host.GroupReceiver("ABC")
+		if r, ok := gotReceiver.(*GroupReceiver); !ok {
+			t.Errorf("expected a hostv2 GroupReceiver; got %v", gotReceiver)
+		} else if r.sub != sub {
+			t.Errorf("unexpected subscriber %v", r.sub)
+		}
+		if err != nil {
+			t.Errorf("expected no error; got %v", err)
+		}
+	})
+	t.Run("Error", func(t *testing.T) {
+		mc := gomock.NewController(t)
+		defer mc.Finish()
+		pubsub := mock.NewMockPubSub(mc)
+		pubsub.EXPECT().Subscribe("ABC").Return(nil, errors.New("FIAL"))
+		host := &HostV2{pubsub: pubsub}
+		gotReceiver, err := host.GroupReceiver("ABC")
+		if gotReceiver != nil {
+			t.Errorf("expected a nil hostv2 GroupReceiver; got %v", gotReceiver)
+		}
+		if err == nil {
+			t.Error("expected an error; got none")
+		}
+	})
+}

p2p/host/hostv2/mock/hostv2_mock.go

@@ -0,0 +1,111 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: hostv2.go
+
+// Package mock_hostv2 is a generated GoMock package.
+package mock_hostv2
+
+import (
+	context "context"
+	gomock "github.com/golang/mock/gomock"
+	go_libp2p_pubsub "github.com/libp2p/go-libp2p-pubsub"
+	reflect "reflect"
+)
+
+// MockPubSub is a mock of PubSub interface
+type MockPubSub struct {
+	ctrl     *gomock.Controller
+	recorder *MockPubSubMockRecorder
+}
+
+// MockPubSubMockRecorder is the mock recorder for MockPubSub
+type MockPubSubMockRecorder struct {
+	mock *MockPubSub
+}
+
+// NewMockPubSub creates a new mock instance
+func NewMockPubSub(ctrl *gomock.Controller) *MockPubSub {
+	mock := &MockPubSub{ctrl: ctrl}
+	mock.recorder = &MockPubSubMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockPubSub) EXPECT() *MockPubSubMockRecorder {
+	return m.recorder
+}
+
+// Publish mocks base method
+func (m *MockPubSub) Publish(topic string, data []byte) error {
+	ret := m.ctrl.Call(m, "Publish", topic, data)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Publish indicates an expected call of Publish
+func (mr *MockPubSubMockRecorder) Publish(topic, data interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Publish", reflect.TypeOf((*MockPubSub)(nil).Publish), topic, data)
+}
+
+// Subscribe mocks base method
+func (m *MockPubSub) Subscribe(topic string, opts ...go_libp2p_pubsub.SubOpt) (*go_libp2p_pubsub.Subscription, error) {
+	varargs := []interface{}{topic}
+	for _, a := range opts {
+		varargs = append(varargs, a)
+	}
+	ret := m.ctrl.Call(m, "Subscribe", varargs...)
+	ret0, _ := ret[0].(*go_libp2p_pubsub.Subscription)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Subscribe indicates an expected call of Subscribe
+func (mr *MockPubSubMockRecorder) Subscribe(topic interface{}, opts ...interface{}) *gomock.Call {
+	varargs := append([]interface{}{topic}, opts...)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Subscribe", reflect.TypeOf((*MockPubSub)(nil).Subscribe), varargs...)
+}
+
+// MockSubscription is a mock of Subscription interface
+type MockSubscription struct {
+	ctrl     *gomock.Controller
+	recorder *MockSubscriptionMockRecorder
+}
+
+// MockSubscriptionMockRecorder is the mock recorder for MockSubscription
+type MockSubscriptionMockRecorder struct {
+	mock *MockSubscription
+}
+
+// NewMockSubscription creates a new mock instance
+func NewMockSubscription(ctrl *gomock.Controller) *MockSubscription {
+	mock := &MockSubscription{ctrl: ctrl}
+	mock.recorder = &MockSubscriptionMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockSubscription) EXPECT() *MockSubscriptionMockRecorder {
+	return m.recorder
+}
+
+// Next mocks base method
+func (m *MockSubscription) Next(ctx context.Context) (*go_libp2p_pubsub.Message, error) {
+	ret := m.ctrl.Call(m, "Next", ctx)
+	ret0, _ := ret[0].(*go_libp2p_pubsub.Message)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Next indicates an expected call of Next
+func (mr *MockSubscriptionMockRecorder) Next(ctx interface{}) *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Next", reflect.TypeOf((*MockSubscription)(nil).Next), ctx)
+}
+
+// Cancel mocks base method
+func (m *MockSubscription) Cancel() {
+	m.ctrl.Call(m, "Cancel")
+}
+
+// Cancel indicates an expected call of Cancel
+func (mr *MockSubscriptionMockRecorder) Cancel() *gomock.Call {
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Cancel", reflect.TypeOf((*MockSubscription)(nil).Cancel))
+}