Fix mock

ChainSafe · Nov 26, 2021 · a5dcfde · a5dcfde
1 parent 993b423
commit a5dcfde
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/dot/network/service_test.go b/dot/network/service_test.go
@@ -73,7 +73,10 @@ func createTestService(t *testing.T, cfg *Config) (srvc *Service) {
 
 	if cfg.TransactionHandler == nil {
 		mocktxhandler := &MockTransactionHandler{}
-		mocktxhandler.On("HandleTransactionMessage", mock.AnythingOfType("*network.TransactionMessage")).Return(true, nil)
+		mocktxhandler.On("HandleTransactionMessage",
+			mock.AnythingOfType("peer.ID"),
+			mock.AnythingOfType("*network.TransactionMessage")).
+			Return(true, nil)
 		mocktxhandler.On("TransactionsCount").Return(0)
 		cfg.TransactionHandler = mocktxhandler
 	}

diff --git a/dot/network/test_helpers.go b/dot/network/test_helpers.go
@@ -8,10 +8,11 @@ import (
 	"io"
 	"math/big"
 
+	"github.com/stretchr/testify/mock"
+
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/common/variadic"
-	"github.com/stretchr/testify/mock"
 
 	libp2pnetwork "github.com/libp2p/go-libp2p-core/network"
 	"github.com/libp2p/go-libp2p-core/peer"
@@ -75,7 +76,7 @@ func NewMockTransactionHandler() *MockTransactionHandler {
 	mocktxhandler.On("HandleTransactionMessage",
 		mock.AnythingOfType("peer.ID"),
 		mock.AnythingOfType("*network.TransactionMessage")).
-		Return(nil)
+		Return(true, nil)
 	mocktxhandler.On("TransactionsCount").Return(0)
 	return mocktxhandler
 }