From 6cd9a60db1e66c69cb50c14091bcba1715bcf117 Mon Sep 17 00:00:00 2001
From: Diego <diego2737@gmail.com>
Date: Wed, 18 Sep 2024 17:18:52 -0300
Subject: [PATCH] Add more tests

---
 dot/network/mock_warp_sync_provider_test.go | 55 ++++++++++++++++++
 dot/network/mocks_generate_test.go          |  1 +
 dot/network/warp_sync_test.go               | 63 +++++++++++++++++++++
 3 files changed, 119 insertions(+)
 create mode 100644 dot/network/mock_warp_sync_provider_test.go

diff --git a/dot/network/mock_warp_sync_provider_test.go b/dot/network/mock_warp_sync_provider_test.go
new file mode 100644
index 0000000000..4b99b3a346
--- /dev/null
+++ b/dot/network/mock_warp_sync_provider_test.go
@@ -0,0 +1,55 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/ChainSafe/gossamer/dot/network (interfaces: WarpSyncProvider)
+//
+// Generated by this command:
+//
+//	mockgen -destination=mock_warp_sync_provider_test.go -package network . WarpSyncProvider
+//
+
+// Package network is a generated GoMock package.
+package network
+
+import (
+	reflect "reflect"
+
+	common "github.com/ChainSafe/gossamer/lib/common"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockWarpSyncProvider is a mock of WarpSyncProvider interface.
+type MockWarpSyncProvider struct {
+	ctrl     *gomock.Controller
+	recorder *MockWarpSyncProviderMockRecorder
+}
+
+// MockWarpSyncProviderMockRecorder is the mock recorder for MockWarpSyncProvider.
+type MockWarpSyncProviderMockRecorder struct {
+	mock *MockWarpSyncProvider
+}
+
+// NewMockWarpSyncProvider creates a new mock instance.
+func NewMockWarpSyncProvider(ctrl *gomock.Controller) *MockWarpSyncProvider {
+	mock := &MockWarpSyncProvider{ctrl: ctrl}
+	mock.recorder = &MockWarpSyncProviderMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockWarpSyncProvider) EXPECT() *MockWarpSyncProviderMockRecorder {
+	return m.recorder
+}
+
+// generate mocks base method.
+func (m *MockWarpSyncProvider) generate(arg0 common.Hash) ([]byte, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "generate", arg0)
+	ret0, _ := ret[0].([]byte)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// generate indicates an expected call of generate.
+func (mr *MockWarpSyncProviderMockRecorder) generate(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "generate", reflect.TypeOf((*MockWarpSyncProvider)(nil).generate), arg0)
+}
diff --git a/dot/network/mocks_generate_test.go b/dot/network/mocks_generate_test.go
index d170388470..fd7f793e78 100644
--- a/dot/network/mocks_generate_test.go
+++ b/dot/network/mocks_generate_test.go
@@ -6,5 +6,6 @@ package network
 //go:generate mockgen -destination=mock_telemetry_test.go -package $GOPACKAGE . Telemetry
 //go:generate mockgen -destination=mock_syncer_test.go -package $GOPACKAGE . Syncer
 //go:generate mockgen -destination=mock_block_state_test.go -package $GOPACKAGE . BlockState
+//go:generate mockgen -destination=mock_warp_sync_provider_test.go -package $GOPACKAGE . WarpSyncProvider
 //go:generate mockgen -destination=mock_transaction_handler_test.go -package $GOPACKAGE . TransactionHandler
 //go:generate mockgen -destination=mock_stream_test.go -package $GOPACKAGE github.com/libp2p/go-libp2p/core/network Stream
diff --git a/dot/network/warp_sync_test.go b/dot/network/warp_sync_test.go
index b961a1c3cc..e80def0345 100644
--- a/dot/network/warp_sync_test.go
+++ b/dot/network/warp_sync_test.go
@@ -4,16 +4,19 @@
 package network
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/ChainSafe/gossamer/dot/network/messages"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/stretchr/testify/require"
+	gomock "go.uber.org/mock/gomock"
 )
 
 func TestDecodeWarpSyncMessage(t *testing.T) {
 	t.Parallel()
+
 	testWarpReqMessage := &messages.WarpProofRequest{
 		Begin: common.EmptyHash,
 	}
@@ -29,3 +32,63 @@ func TestDecodeWarpSyncMessage(t *testing.T) {
 	require.True(t, ok)
 	require.Equal(t, testWarpReqMessage, req)
 }
+
+func createServiceWithWarpSyncHelper(t *testing.T, warpSyncProvider WarpSyncProvider) *Service {
+	t.Helper()
+
+	config := &Config{
+		BasePath:         t.TempDir(),
+		Port:             availablePort(t),
+		NoBootstrap:      true,
+		NoMDNS:           true,
+		warpSyncProvider: warpSyncProvider,
+	}
+
+	srvc := createTestService(t, config)
+	srvc.noGossip = true
+	handler := newTestStreamHandler(decodeSyncMessage)
+	srvc.host.registerStreamHandler(srvc.host.protocolID, handler.handleStream)
+
+	return srvc
+}
+
+func TestHandleWarpSyncRequestOk(t *testing.T) {
+	t.Parallel()
+
+	expectedProof := []byte{0x01}
+
+	ctrl := gomock.NewController(t)
+
+	warpSyncProvider := NewMockWarpSyncProvider(ctrl)
+	warpSyncProvider.EXPECT().generate(common.EmptyHash).Return(expectedProof, nil).Times(1)
+
+	srvc := createServiceWithWarpSyncHelper(t, warpSyncProvider)
+
+	req := messages.WarpProofRequest{
+		Begin: common.EmptyHash,
+	}
+
+	resp, err := srvc.handleWarpSyncRequest(req)
+	require.NoError(t, err)
+	require.Equal(t, expectedProof, resp)
+}
+
+func TestHandleWarpSyncRequestError(t *testing.T) {
+	t.Parallel()
+
+	expectedError := fmt.Errorf("error generating proof")
+	ctrl := gomock.NewController(t)
+
+	warpSyncProvider := NewMockWarpSyncProvider(ctrl)
+	warpSyncProvider.EXPECT().generate(common.EmptyHash).Return(nil, expectedError).Times(1)
+
+	srvc := createServiceWithWarpSyncHelper(t, warpSyncProvider)
+
+	req := messages.WarpProofRequest{
+		Begin: common.EmptyHash,
+	}
+
+	resp, err := srvc.handleWarpSyncRequest(req)
+	require.Nil(t, resp)
+	require.ErrorIs(t, err, expectedError)
+}