diff --git a/dot/core/helpers_test.go b/dot/core/helpers_test.go
index f4f4c7df87..92b538b154 100644
--- a/dot/core/helpers_test.go
+++ b/dot/core/helpers_test.go
@@ -7,8 +7,10 @@ import (
 	"path/filepath"
 	"testing"
 
-	coremocks "github.com/ChainSafe/gossamer/dot/core/mocks"
+	"github.com/ChainSafe/gossamer/dot/network"
+	"github.com/ChainSafe/gossamer/dot/peerset"
 	"github.com/ChainSafe/gossamer/dot/state"
+	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/internal/log"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/crypto/sr25519"
@@ -19,21 +21,23 @@ import (
 	"github.com/ChainSafe/gossamer/lib/runtime/wasmer"
 	"github.com/ChainSafe/gossamer/lib/utils"
 	"github.com/golang/mock/gomock"
-	"github.com/stretchr/testify/mock"
+	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/stretchr/testify/require"
 )
 
 // NewTestService creates a new test core service
 func NewTestService(t *testing.T, cfg *Config) *Service {
 	t.Helper()
+	ctrl := gomock.NewController(t)
 
 	if cfg == nil {
 		cfg = &Config{}
 	}
 
 	if cfg.DigestHandler == nil {
-		cfg.DigestHandler = new(coremocks.DigestHandler)
-		cfg.DigestHandler.(*coremocks.DigestHandler).On("HandleDigests", mock.AnythingOfType("*types.Header"))
+		digestHandler := NewMockDigestHandler(ctrl)
+		digestHandler.EXPECT().HandleDigests(gomock.AssignableToTypeOf(new(types.Header)))
+		cfg.DigestHandler = digestHandler
 	}
 
 	if cfg.Keystore == nil {
@@ -56,7 +60,6 @@ func NewTestService(t *testing.T, cfg *Config) *Service {
 	if cfg.BlockState == nil || cfg.StorageState == nil ||
 		cfg.TransactionState == nil || cfg.EpochState == nil ||
 		cfg.CodeSubstitutedState == nil {
-		ctrl := gomock.NewController(t)
 		telemetryMock := NewMockClient(ctrl)
 		telemetryMock.EXPECT().SendMessage(gomock.Any()).AnyTimes()
 
@@ -123,10 +126,13 @@ func NewTestService(t *testing.T, cfg *Config) *Service {
 	cfg.BlockState.StoreRuntime(cfg.BlockState.BestBlockHash(), cfg.Runtime)
 
 	if cfg.Network == nil {
-		net := new(coremocks.Network)
-		net.On("GossipMessage", mock.AnythingOfType("*network.TransactionMessage"))
-		net.On("IsSynced").Return(true)
-		net.On("ReportPeer", mock.AnythingOfType("peerset.ReputationChange"), mock.AnythingOfType("peer.ID"))
+		net := NewMockNetwork(ctrl)
+		net.EXPECT().GossipMessage(gomock.AssignableToTypeOf(new(network.TransactionMessage)))
+		net.EXPECT().IsSynced().Return(true)
+		net.EXPECT().ReportPeer(
+			gomock.AssignableToTypeOf(peerset.ReputationChange{}),
+			gomock.AssignableToTypeOf(peer.ID("")),
+		)
 		cfg.Network = net
 	}
 
diff --git a/dot/core/messages_integration_test.go b/dot/core/messages_integration_test.go
index 5fca745e56..f95c083559 100644
--- a/dot/core/messages_integration_test.go
+++ b/dot/core/messages_integration_test.go
@@ -16,7 +16,6 @@ import (
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
-	"github.com/ChainSafe/gossamer/dot/core/mocks"
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/sync"
@@ -70,7 +69,9 @@ func createExtrinsic(t *testing.T, rt runtime.Instance, genHash common.Hash, non
 }
 
 func TestService_HandleBlockProduced(t *testing.T) {
-	net := new(mocks.Network)
+	ctrl := gomock.NewController(t)
+
+	net := NewMockNetwork(ctrl)
 	cfg := &Config{
 		Network:  net,
 		Keystore: keystore.NewGlobalKeystore(),
@@ -105,7 +106,7 @@ func TestService_HandleBlockProduced(t *testing.T) {
 		BestBlock:      true,
 	}
 
-	net.On("GossipMessage", expected)
+	net.EXPECT().GossipMessage(expected)
 
 	state, err := s.storageState.TrieState(nil)
 	require.NoError(t, err)
@@ -114,7 +115,6 @@ func TestService_HandleBlockProduced(t *testing.T) {
 	require.NoError(t, err)
 
 	time.Sleep(time.Second)
-	net.AssertCalled(t, "GossipMessage", expected)
 }
 
 func TestService_HandleTransactionMessage(t *testing.T) {
diff --git a/dot/core/mocks/block_state.go b/dot/core/mocks/block_state.go
deleted file mode 100644
index 43f779983c..0000000000
--- a/dot/core/mocks/block_state.go
+++ /dev/null
@@ -1,442 +0,0 @@
-// Code generated by mockery v2.9.4. DO NOT EDIT.
-
-package mocks
-
-import (
-	big "math/big"
-
-	common "github.com/ChainSafe/gossamer/lib/common"
-
-	mock "github.com/stretchr/testify/mock"
-
-	runtime "github.com/ChainSafe/gossamer/lib/runtime"
-
-	storage "github.com/ChainSafe/gossamer/lib/runtime/storage"
-
-	types "github.com/ChainSafe/gossamer/dot/types"
-)
-
-// BlockState is an autogenerated mock type for the BlockState type
-type BlockState struct {
-	mock.Mock
-}
-
-// AddBlock provides a mock function with given fields: _a0
-func (_m *BlockState) AddBlock(_a0 *types.Block) error {
-	ret := _m.Called(_a0)
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*types.Block) error); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// BestBlock provides a mock function with given fields:
-func (_m *BlockState) BestBlock() (*types.Block, error) {
-	ret := _m.Called()
-
-	var r0 *types.Block
-	if rf, ok := ret.Get(0).(func() *types.Block); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*types.Block)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// BestBlockHash provides a mock function with given fields:
-func (_m *BlockState) BestBlockHash() common.Hash {
-	ret := _m.Called()
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func() common.Hash); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	return r0
-}
-
-// BestBlockHeader provides a mock function with given fields:
-func (_m *BlockState) BestBlockHeader() (*types.Header, error) {
-	ret := _m.Called()
-
-	var r0 *types.Header
-	if rf, ok := ret.Get(0).(func() *types.Header); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*types.Header)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// BestBlockNumber provides a mock function with given fields:
-func (_m *BlockState) BestBlockNumber() (*big.Int, error) {
-	ret := _m.Called()
-
-	var r0 *big.Int
-	if rf, ok := ret.Get(0).(func() *big.Int); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*big.Int)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// BestBlockStateRoot provides a mock function with given fields:
-func (_m *BlockState) BestBlockStateRoot() (common.Hash, error) {
-	ret := _m.Called()
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func() common.Hash); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// FreeFinalisedNotifierChannel provides a mock function with given fields: ch
-func (_m *BlockState) FreeFinalisedNotifierChannel(ch chan *types.FinalisationInfo) {
-	_m.Called(ch)
-}
-
-// FreeImportedBlockNotifierChannel provides a mock function with given fields: ch
-func (_m *BlockState) FreeImportedBlockNotifierChannel(ch chan *types.Block) {
-	_m.Called(ch)
-}
-
-// GenesisHash provides a mock function with given fields:
-func (_m *BlockState) GenesisHash() common.Hash {
-	ret := _m.Called()
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func() common.Hash); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	return r0
-}
-
-// GetAllBlocksAtDepth provides a mock function with given fields: hash
-func (_m *BlockState) GetAllBlocksAtDepth(hash common.Hash) []common.Hash {
-	ret := _m.Called(hash)
-
-	var r0 []common.Hash
-	if rf, ok := ret.Get(0).(func(common.Hash) []common.Hash); ok {
-		r0 = rf(hash)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([]common.Hash)
-		}
-	}
-
-	return r0
-}
-
-// GetBlockBody provides a mock function with given fields: hash
-func (_m *BlockState) GetBlockBody(hash common.Hash) (*types.Body, error) {
-	ret := _m.Called(hash)
-
-	var r0 *types.Body
-	if rf, ok := ret.Get(0).(func(common.Hash) *types.Body); ok {
-		r0 = rf(hash)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*types.Body)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash) error); ok {
-		r1 = rf(hash)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetBlockByHash provides a mock function with given fields: _a0
-func (_m *BlockState) GetBlockByHash(_a0 common.Hash) (*types.Block, error) {
-	ret := _m.Called(_a0)
-
-	var r0 *types.Block
-	if rf, ok := ret.Get(0).(func(common.Hash) *types.Block); ok {
-		r0 = rf(_a0)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*types.Block)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash) error); ok {
-		r1 = rf(_a0)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetBlockStateRoot provides a mock function with given fields: bhash
-func (_m *BlockState) GetBlockStateRoot(bhash common.Hash) (common.Hash, error) {
-	ret := _m.Called(bhash)
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func(common.Hash) common.Hash); ok {
-		r0 = rf(bhash)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash) error); ok {
-		r1 = rf(bhash)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetFinalisedHash provides a mock function with given fields: _a0, _a1
-func (_m *BlockState) GetFinalisedHash(_a0 uint64, _a1 uint64) (common.Hash, error) {
-	ret := _m.Called(_a0, _a1)
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func(uint64, uint64) common.Hash); ok {
-		r0 = rf(_a0, _a1)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(uint64, uint64) error); ok {
-		r1 = rf(_a0, _a1)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetFinalisedHeader provides a mock function with given fields: _a0, _a1
-func (_m *BlockState) GetFinalisedHeader(_a0 uint64, _a1 uint64) (*types.Header, error) {
-	ret := _m.Called(_a0, _a1)
-
-	var r0 *types.Header
-	if rf, ok := ret.Get(0).(func(uint64, uint64) *types.Header); ok {
-		r0 = rf(_a0, _a1)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*types.Header)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(uint64, uint64) error); ok {
-		r1 = rf(_a0, _a1)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetFinalisedNotifierChannel provides a mock function with given fields:
-func (_m *BlockState) GetFinalisedNotifierChannel() chan *types.FinalisationInfo {
-	ret := _m.Called()
-
-	var r0 chan *types.FinalisationInfo
-	if rf, ok := ret.Get(0).(func() chan *types.FinalisationInfo); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(chan *types.FinalisationInfo)
-		}
-	}
-
-	return r0
-}
-
-// GetImportedBlockNotifierChannel provides a mock function with given fields:
-func (_m *BlockState) GetImportedBlockNotifierChannel() chan *types.Block {
-	ret := _m.Called()
-
-	var r0 chan *types.Block
-	if rf, ok := ret.Get(0).(func() chan *types.Block); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(chan *types.Block)
-		}
-	}
-
-	return r0
-}
-
-// GetRuntime provides a mock function with given fields: _a0
-func (_m *BlockState) GetRuntime(_a0 *common.Hash) (runtime.Instance, error) {
-	ret := _m.Called(_a0)
-
-	var r0 runtime.Instance
-	if rf, ok := ret.Get(0).(func(*common.Hash) runtime.Instance); ok {
-		r0 = rf(_a0)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(runtime.Instance)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash) error); ok {
-		r1 = rf(_a0)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetSlotForBlock provides a mock function with given fields: _a0
-func (_m *BlockState) GetSlotForBlock(_a0 common.Hash) (uint64, error) {
-	ret := _m.Called(_a0)
-
-	var r0 uint64
-	if rf, ok := ret.Get(0).(func(common.Hash) uint64); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Get(0).(uint64)
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash) error); ok {
-		r1 = rf(_a0)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// HandleRuntimeChanges provides a mock function with given fields: newState, in, bHash
-func (_m *BlockState) HandleRuntimeChanges(newState *storage.TrieState, in runtime.Instance, bHash common.Hash) error {
-	ret := _m.Called(newState, in, bHash)
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*storage.TrieState, runtime.Instance, common.Hash) error); ok {
-		r0 = rf(newState, in, bHash)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// HighestCommonAncestor provides a mock function with given fields: a, b
-func (_m *BlockState) HighestCommonAncestor(a common.Hash, b common.Hash) (common.Hash, error) {
-	ret := _m.Called(a, b)
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func(common.Hash, common.Hash) common.Hash); ok {
-		r0 = rf(a, b)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash, common.Hash) error); ok {
-		r1 = rf(a, b)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// StoreRuntime provides a mock function with given fields: _a0, _a1
-func (_m *BlockState) StoreRuntime(_a0 common.Hash, _a1 runtime.Instance) {
-	_m.Called(_a0, _a1)
-}
-
-// SubChain provides a mock function with given fields: start, end
-func (_m *BlockState) SubChain(start common.Hash, end common.Hash) ([]common.Hash, error) {
-	ret := _m.Called(start, end)
-
-	var r0 []common.Hash
-	if rf, ok := ret.Get(0).(func(common.Hash, common.Hash) []common.Hash); ok {
-		r0 = rf(start, end)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([]common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash, common.Hash) error); ok {
-		r1 = rf(start, end)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
diff --git a/dot/core/mocks/digest_handler.go b/dot/core/mocks/digest_handler.go
deleted file mode 100644
index c6f23a4c3c..0000000000
--- a/dot/core/mocks/digest_handler.go
+++ /dev/null
@@ -1,18 +0,0 @@
-// Code generated by mockery v2.9.4. DO NOT EDIT.
-
-package mocks
-
-import (
-	types "github.com/ChainSafe/gossamer/dot/types"
-	mock "github.com/stretchr/testify/mock"
-)
-
-// DigestHandler is an autogenerated mock type for the DigestHandler type
-type DigestHandler struct {
-	mock.Mock
-}
-
-// HandleDigests provides a mock function with given fields: header
-func (_m *DigestHandler) HandleDigests(header *types.Header) {
-	_m.Called(header)
-}
diff --git a/dot/core/mocks/network.go b/dot/core/mocks/network.go
deleted file mode 100644
index 37cd0c072f..0000000000
--- a/dot/core/mocks/network.go
+++ /dev/null
@@ -1,41 +0,0 @@
-// Code generated by mockery v2.9.4. DO NOT EDIT.
-
-package mocks
-
-import (
-	network "github.com/ChainSafe/gossamer/dot/network"
-	mock "github.com/stretchr/testify/mock"
-
-	peer "github.com/libp2p/go-libp2p-core/peer"
-
-	peerset "github.com/ChainSafe/gossamer/dot/peerset"
-)
-
-// Network is an autogenerated mock type for the Network type
-type Network struct {
-	mock.Mock
-}
-
-// GossipMessage provides a mock function with given fields: _a0
-func (_m *Network) GossipMessage(_a0 network.NotificationsMessage) {
-	_m.Called(_a0)
-}
-
-// IsSynced provides a mock function with given fields:
-func (_m *Network) IsSynced() bool {
-	ret := _m.Called()
-
-	var r0 bool
-	if rf, ok := ret.Get(0).(func() bool); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Get(0).(bool)
-	}
-
-	return r0
-}
-
-// ReportPeer provides a mock function with given fields: change, p
-func (_m *Network) ReportPeer(change peerset.ReputationChange, p peer.ID) {
-	_m.Called(change, p)
-}
diff --git a/dot/core/mocks/storage_state.go b/dot/core/mocks/storage_state.go
deleted file mode 100644
index 320a6165d0..0000000000
--- a/dot/core/mocks/storage_state.go
+++ /dev/null
@@ -1,180 +0,0 @@
-// Code generated by mockery v2.9.4. DO NOT EDIT.
-
-package mocks
-
-import (
-	common "github.com/ChainSafe/gossamer/lib/common"
-
-	mock "github.com/stretchr/testify/mock"
-
-	storage "github.com/ChainSafe/gossamer/lib/runtime/storage"
-
-	types "github.com/ChainSafe/gossamer/dot/types"
-)
-
-// StorageState is an autogenerated mock type for the StorageState type
-type StorageState struct {
-	mock.Mock
-}
-
-// GenerateTrieProof provides a mock function with given fields: stateRoot, keys
-func (_m *StorageState) GenerateTrieProof(stateRoot common.Hash, keys [][]byte) ([][]byte, error) {
-	ret := _m.Called(stateRoot, keys)
-
-	var r0 [][]byte
-	if rf, ok := ret.Get(0).(func(common.Hash, [][]byte) [][]byte); ok {
-		r0 = rf(stateRoot, keys)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([][]byte)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(common.Hash, [][]byte) error); ok {
-		r1 = rf(stateRoot, keys)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetStateRootFromBlock provides a mock function with given fields: bhash
-func (_m *StorageState) GetStateRootFromBlock(bhash *common.Hash) (*common.Hash, error) {
-	ret := _m.Called(bhash)
-
-	var r0 *common.Hash
-	if rf, ok := ret.Get(0).(func(*common.Hash) *common.Hash); ok {
-		r0 = rf(bhash)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash) error); ok {
-		r1 = rf(bhash)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// GetStorage provides a mock function with given fields: root, key
-func (_m *StorageState) GetStorage(root *common.Hash, key []byte) ([]byte, error) {
-	ret := _m.Called(root, key)
-
-	var r0 []byte
-	if rf, ok := ret.Get(0).(func(*common.Hash, []byte) []byte); ok {
-		r0 = rf(root, key)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([]byte)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash, []byte) error); ok {
-		r1 = rf(root, key)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// LoadCode provides a mock function with given fields: root
-func (_m *StorageState) LoadCode(root *common.Hash) ([]byte, error) {
-	ret := _m.Called(root)
-
-	var r0 []byte
-	if rf, ok := ret.Get(0).(func(*common.Hash) []byte); ok {
-		r0 = rf(root)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([]byte)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash) error); ok {
-		r1 = rf(root)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// LoadCodeHash provides a mock function with given fields: root
-func (_m *StorageState) LoadCodeHash(root *common.Hash) (common.Hash, error) {
-	ret := _m.Called(root)
-
-	var r0 common.Hash
-	if rf, ok := ret.Get(0).(func(*common.Hash) common.Hash); ok {
-		r0 = rf(root)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(common.Hash)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash) error); ok {
-		r1 = rf(root)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// Lock provides a mock function with given fields:
-func (_m *StorageState) Lock() {
-	_m.Called()
-}
-
-// StoreTrie provides a mock function with given fields: _a0, _a1
-func (_m *StorageState) StoreTrie(_a0 *storage.TrieState, _a1 *types.Header) error {
-	ret := _m.Called(_a0, _a1)
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*storage.TrieState, *types.Header) error); ok {
-		r0 = rf(_a0, _a1)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// TrieState provides a mock function with given fields: root
-func (_m *StorageState) TrieState(root *common.Hash) (*storage.TrieState, error) {
-	ret := _m.Called(root)
-
-	var r0 *storage.TrieState
-	if rf, ok := ret.Get(0).(func(*common.Hash) *storage.TrieState); ok {
-		r0 = rf(root)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*storage.TrieState)
-		}
-	}
-
-	var r1 error
-	if rf, ok := ret.Get(1).(func(*common.Hash) error); ok {
-		r1 = rf(root)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// Unlock provides a mock function with given fields:
-func (_m *StorageState) Unlock() {
-	_m.Called()
-}
diff --git a/dot/core/service_integration_test.go b/dot/core/service_integration_test.go
index 7ecd5be0a2..22126e23f9 100644
--- a/dot/core/service_integration_test.go
+++ b/dot/core/service_integration_test.go
@@ -15,7 +15,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/ChainSafe/gossamer/dot/core/mocks"
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/sync"
@@ -63,7 +62,8 @@ func TestStartService(t *testing.T) {
 }
 
 func TestAnnounceBlock(t *testing.T) {
-	net := new(mocks.Network)
+	ctrl := gomock.NewController(t)
+	net := NewMockNetwork(ctrl)
 	cfg := &Config{
 		Network: net,
 	}
@@ -98,7 +98,7 @@ func TestAnnounceBlock(t *testing.T) {
 		BestBlock:      true,
 	}
 
-	net.On("GossipMessage", expected)
+	net.EXPECT().GossipMessage(expected)
 
 	state, err := s.storageState.TrieState(nil)
 	require.NoError(t, err)
@@ -107,7 +107,6 @@ func TestAnnounceBlock(t *testing.T) {
 	require.NoError(t, err)
 
 	time.Sleep(time.Second)
-	net.AssertCalled(t, "GossipMessage", expected)
 }
 
 func TestService_InsertKey(t *testing.T) {
@@ -839,18 +838,20 @@ func createNewBlockAndStoreDataAtBlock(t *testing.T, s *Service,
 }
 
 func TestDecodeSessionKeys(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
 	mockInstance := new(runtimemocks.Instance)
 	mockInstance.On("DecodeSessionKeys", mock.AnythingOfType("[]uint8")).Return([]byte{}, nil).Once()
 
-	mockBlockState := new(mocks.BlockState)
-	mockBlockState.On("GetRuntime", mock.AnythingOfType("*common.Hash")).Return(mockInstance, nil).Once()
+	mockBlockState := NewMockBlockState(ctrl)
+	mockBlockState.EXPECT().GetRuntime(gomock.AssignableToTypeOf(new(common.Hash))).
+		Return(mockInstance, nil)
 
 	coreservice := new(Service)
 	coreservice.blockState = mockBlockState
 
 	b, err := coreservice.DecodeSessionKeys([]byte{})
 
-	mockBlockState.AssertCalled(t, "GetRuntime", mock.AnythingOfType("*common.Hash"))
 	mockInstance.AssertCalled(t, "DecodeSessionKeys", []uint8{})
 
 	require.NoError(t, err)
@@ -858,15 +859,17 @@ func TestDecodeSessionKeys(t *testing.T) {
 }
 
 func TestDecodeSessionKeys_WhenGetRuntimeReturnError(t *testing.T) {
-	mockBlockState := new(mocks.BlockState)
-	mockBlockState.On("GetRuntime", mock.AnythingOfType("*common.Hash")).Return(nil, errors.New("problems")).Once()
+	ctrl := gomock.NewController(t)
+
+	mockBlockState := NewMockBlockState(ctrl)
+	mockBlockState.EXPECT().GetRuntime(gomock.AssignableToTypeOf(new(common.Hash))).
+		Return(nil, errors.New("problems"))
 
 	coreservice := new(Service)
 	coreservice.blockState = mockBlockState
 
 	b, err := coreservice.DecodeSessionKeys([]byte{})
 
-	mockBlockState.AssertCalled(t, "GetRuntime", mock.AnythingOfType("*common.Hash"))
 	require.Error(t, err, "problems")
 	require.Nil(t, b)
 }
@@ -876,16 +879,18 @@ func TestGetReadProofAt(t *testing.T) {
 	mockedProofs := [][]byte{[]byte("proof01"), []byte("proof02")}
 
 	t.Run("When Has Block Is Empty", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
 		mockedStateRootHash := common.NewHash([]byte("state root hash"))
 		expectedBlockHash := common.NewHash([]byte("expected block hash"))
 
-		mockBlockState := new(mocks.BlockState)
-		mockBlockState.On("BestBlockHash").Return(expectedBlockHash)
-		mockBlockState.On("GetBlockStateRoot", expectedBlockHash).
+		mockBlockState := NewMockBlockState(ctrl)
+		mockBlockState.EXPECT().BestBlockHash().Return(expectedBlockHash)
+		mockBlockState.EXPECT().GetBlockStateRoot(expectedBlockHash).
 			Return(mockedStateRootHash, nil)
 
-		mockStorageStage := new(mocks.StorageState)
-		mockStorageStage.On("GenerateTrieProof", mockedStateRootHash, keysToProof).
+		mockStorageStage := NewMockStorageState(ctrl)
+		mockStorageStage.EXPECT().GenerateTrieProof(mockedStateRootHash, keysToProof).
 			Return(mockedProofs, nil)
 
 		s := &Service{
@@ -897,17 +902,15 @@ func TestGetReadProofAt(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, p, mockedProofs)
 		require.Equal(t, expectedBlockHash, b)
-
-		mockBlockState.AssertCalled(t, "BestBlockHash")
-		mockBlockState.AssertCalled(t, "GetBlockStateRoot", expectedBlockHash)
-		mockStorageStage.AssertCalled(t, "GenerateTrieProof", mockedStateRootHash, keysToProof)
 	})
 
 	t.Run("When GetStateRoot fails", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
 		mockedBlockHash := common.NewHash([]byte("fake block hash"))
 
-		mockBlockState := new(mocks.BlockState)
-		mockBlockState.On("GetBlockStateRoot", mockedBlockHash).
+		mockBlockState := NewMockBlockState(ctrl)
+		mockBlockState.EXPECT().GetBlockStateRoot(mockedBlockHash).
 			Return(common.Hash{}, errors.New("problems while getting state root"))
 
 		s := &Service{
@@ -921,15 +924,17 @@ func TestGetReadProofAt(t *testing.T) {
 	})
 
 	t.Run("When GenerateTrieProof fails", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
 		mockedBlockHash := common.NewHash([]byte("fake block hash"))
 		mockedStateRootHash := common.NewHash([]byte("state root hash"))
 
-		mockBlockState := new(mocks.BlockState)
-		mockBlockState.On("GetBlockStateRoot", mockedBlockHash).
-			Return(mockedStateRootHash, nil)
+		mockBlockState := NewMockBlockState(ctrl)
+		mockBlockState.EXPECT().GetBlockStateRoot(mockedBlockHash).
+			Return(common.Hash{}, errors.New("problems while getting state root"))
 
-		mockStorageStage := new(mocks.StorageState)
-		mockStorageStage.On("GenerateTrieProof", mockedStateRootHash, keysToProof).
+		mockStorageStage := NewMockStorageState(ctrl)
+		mockStorageStage.EXPECT().GenerateTrieProof(mockedStateRootHash, keysToProof).
 			Return(nil, errors.New("problems to generate trie proof"))
 
 		s := &Service{
diff --git a/dot/rpc/http_test.go b/dot/rpc/http_test.go
index a5bbb29598..800846e0d6 100644
--- a/dot/rpc/http_test.go
+++ b/dot/rpc/http_test.go
@@ -14,9 +14,11 @@ import (
 	"time"
 
 	rtstorage "github.com/ChainSafe/gossamer/lib/runtime/storage"
+	"github.com/libp2p/go-libp2p-core/peer"
 
 	"github.com/ChainSafe/gossamer/dot/core"
-	coremocks "github.com/ChainSafe/gossamer/dot/core/mocks"
+	"github.com/ChainSafe/gossamer/dot/network"
+	"github.com/ChainSafe/gossamer/dot/peerset"
 	"github.com/ChainSafe/gossamer/dot/rpc/modules"
 	"github.com/ChainSafe/gossamer/dot/rpc/modules/mocks"
 	"github.com/ChainSafe/gossamer/dot/state"
@@ -364,6 +366,8 @@ func externalIP() (string, error) {
 }
 
 //go:generate mockgen -destination=mock_telemetry_test.go -package $GOPACKAGE github.com/ChainSafe/gossamer/dot/telemetry Client
+//go:generate mockgen -destination=mock_digesthandler_test.go -package $GOPACKAGE github.com/ChainSafe/gossamer/dot/core DigestHandler
+//go:generate mockgen -destination=mock_network_test.go -package $GOPACKAGE github.com/ChainSafe/gossamer/dot/core Network
 
 func newCoreServiceTest(t *testing.T) *core.Service {
 	t.Helper()
@@ -403,8 +407,11 @@ func newCoreServiceTest(t *testing.T) *core.Service {
 		CodeSubstitutedState: stateSrvc.Base,
 	}
 
-	cfg.DigestHandler = new(coremocks.DigestHandler)
-	cfg.DigestHandler.(*coremocks.DigestHandler).On("HandleDigests", mock.AnythingOfType("*types.Header"))
+	digestHandler := NewMockDigestHandler(ctrl)
+	digestHandler.EXPECT().
+		HandleDigests(gomock.AssignableToTypeOf(new(types.Header))).
+		AnyTimes()
+	cfg.DigestHandler = digestHandler
 
 	cfg.Keystore = keystore.NewGlobalKeystore()
 	kp, err := sr25519.GenerateKeypair()
@@ -431,11 +438,15 @@ func newCoreServiceTest(t *testing.T) *core.Service {
 
 	cfg.BlockState.StoreRuntime(cfg.BlockState.BestBlockHash(), cfg.Runtime)
 
-	net := new(coremocks.Network)
-	net.On("GossipMessage", mock.AnythingOfType("*network.TransactionMessage"))
-	net.On("IsSynced").Return(true)
-	net.On("ReportPeer", mock.AnythingOfType("peerset.ReputationChange"), mock.AnythingOfType("peer.ID"))
-
+	net := NewMockNetwork(ctrl)
+	net.EXPECT().
+		GossipMessage(gomock.AssignableToTypeOf(new(network.TransactionMessage))).
+		AnyTimes()
+	net.EXPECT().IsSynced().Return(true).AnyTimes()
+	net.EXPECT().ReportPeer(
+		gomock.AssignableToTypeOf(new(peerset.ReputationChange)),
+		gomock.AssignableToTypeOf(peer.ID(""))).
+		AnyTimes()
 	cfg.Network = net
 
 	cfg.CodeSubstitutes = make(map[common.Hash]string)
diff --git a/dot/rpc/mock_digesthandler_test.go b/dot/rpc/mock_digesthandler_test.go
new file mode 100644
index 0000000000..947c0257f8
--- /dev/null
+++ b/dot/rpc/mock_digesthandler_test.go
@@ -0,0 +1,47 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/ChainSafe/gossamer/dot/core (interfaces: DigestHandler)
+
+// Package rpc is a generated GoMock package.
+package rpc
+
+import (
+	reflect "reflect"
+
+	types "github.com/ChainSafe/gossamer/dot/types"
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockDigestHandler is a mock of DigestHandler interface.
+type MockDigestHandler struct {
+	ctrl     *gomock.Controller
+	recorder *MockDigestHandlerMockRecorder
+}
+
+// MockDigestHandlerMockRecorder is the mock recorder for MockDigestHandler.
+type MockDigestHandlerMockRecorder struct {
+	mock *MockDigestHandler
+}
+
+// NewMockDigestHandler creates a new mock instance.
+func NewMockDigestHandler(ctrl *gomock.Controller) *MockDigestHandler {
+	mock := &MockDigestHandler{ctrl: ctrl}
+	mock.recorder = &MockDigestHandlerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockDigestHandler) EXPECT() *MockDigestHandlerMockRecorder {
+	return m.recorder
+}
+
+// HandleDigests mocks base method.
+func (m *MockDigestHandler) HandleDigests(arg0 *types.Header) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "HandleDigests", arg0)
+}
+
+// HandleDigests indicates an expected call of HandleDigests.
+func (mr *MockDigestHandlerMockRecorder) HandleDigests(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "HandleDigests", reflect.TypeOf((*MockDigestHandler)(nil).HandleDigests), arg0)
+}
diff --git a/dot/rpc/mock_network_test.go b/dot/rpc/mock_network_test.go
new file mode 100644
index 0000000000..f2809eb44a
--- /dev/null
+++ b/dot/rpc/mock_network_test.go
@@ -0,0 +1,75 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/ChainSafe/gossamer/dot/core (interfaces: Network)
+
+// Package rpc is a generated GoMock package.
+package rpc
+
+import (
+	reflect "reflect"
+
+	network "github.com/ChainSafe/gossamer/dot/network"
+	peerset "github.com/ChainSafe/gossamer/dot/peerset"
+	gomock "github.com/golang/mock/gomock"
+	peer "github.com/libp2p/go-libp2p-core/peer"
+)
+
+// MockNetwork is a mock of Network interface.
+type MockNetwork struct {
+	ctrl     *gomock.Controller
+	recorder *MockNetworkMockRecorder
+}
+
+// MockNetworkMockRecorder is the mock recorder for MockNetwork.
+type MockNetworkMockRecorder struct {
+	mock *MockNetwork
+}
+
+// NewMockNetwork creates a new mock instance.
+func NewMockNetwork(ctrl *gomock.Controller) *MockNetwork {
+	mock := &MockNetwork{ctrl: ctrl}
+	mock.recorder = &MockNetworkMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockNetwork) EXPECT() *MockNetworkMockRecorder {
+	return m.recorder
+}
+
+// GossipMessage mocks base method.
+func (m *MockNetwork) GossipMessage(arg0 network.NotificationsMessage) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "GossipMessage", arg0)
+}
+
+// GossipMessage indicates an expected call of GossipMessage.
+func (mr *MockNetworkMockRecorder) GossipMessage(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GossipMessage", reflect.TypeOf((*MockNetwork)(nil).GossipMessage), arg0)
+}
+
+// IsSynced mocks base method.
+func (m *MockNetwork) IsSynced() bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "IsSynced")
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// IsSynced indicates an expected call of IsSynced.
+func (mr *MockNetworkMockRecorder) IsSynced() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsSynced", reflect.TypeOf((*MockNetwork)(nil).IsSynced))
+}
+
+// ReportPeer mocks base method.
+func (m *MockNetwork) ReportPeer(arg0 peerset.ReputationChange, arg1 peer.ID) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "ReportPeer", arg0, arg1)
+}
+
+// ReportPeer indicates an expected call of ReportPeer.
+func (mr *MockNetworkMockRecorder) ReportPeer(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReportPeer", reflect.TypeOf((*MockNetwork)(nil).ReportPeer), arg0, arg1)
+}
diff --git a/dot/rpc/modules/api.go b/dot/rpc/modules/api.go
index b0f2e4488b..bff31c410e 100644
--- a/dot/rpc/modules/api.go
+++ b/dot/rpc/modules/api.go
@@ -72,6 +72,8 @@ type NetworkAPI interface {
 	RemoveReservedPeers(addrs ...string) error
 }
 
+//go:generate mockery --name BlockProducerAPI --structname BlockProducerAPI --case underscore --keeptree
+
 // BlockProducerAPI is the interface for BlockProducer methods
 type BlockProducerAPI interface {
 	Pause() error
diff --git a/dot/rpc/modules/mock_network_test.go b/dot/rpc/modules/mock_network_test.go
new file mode 100644
index 0000000000..cc059d9d50
--- /dev/null
+++ b/dot/rpc/modules/mock_network_test.go
@@ -0,0 +1,75 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/ChainSafe/gossamer/dot/core (interfaces: Network)
+
+// Package modules is a generated GoMock package.
+package modules
+
+import (
+	reflect "reflect"
+
+	network "github.com/ChainSafe/gossamer/dot/network"
+	peerset "github.com/ChainSafe/gossamer/dot/peerset"
+	gomock "github.com/golang/mock/gomock"
+	peer "github.com/libp2p/go-libp2p-core/peer"
+)
+
+// MockNetwork is a mock of Network interface.
+type MockNetwork struct {
+	ctrl     *gomock.Controller
+	recorder *MockNetworkMockRecorder
+}
+
+// MockNetworkMockRecorder is the mock recorder for MockNetwork.
+type MockNetworkMockRecorder struct {
+	mock *MockNetwork
+}
+
+// NewMockNetwork creates a new mock instance.
+func NewMockNetwork(ctrl *gomock.Controller) *MockNetwork {
+	mock := &MockNetwork{ctrl: ctrl}
+	mock.recorder = &MockNetworkMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockNetwork) EXPECT() *MockNetworkMockRecorder {
+	return m.recorder
+}
+
+// GossipMessage mocks base method.
+func (m *MockNetwork) GossipMessage(arg0 network.NotificationsMessage) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "GossipMessage", arg0)
+}
+
+// GossipMessage indicates an expected call of GossipMessage.
+func (mr *MockNetworkMockRecorder) GossipMessage(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GossipMessage", reflect.TypeOf((*MockNetwork)(nil).GossipMessage), arg0)
+}
+
+// IsSynced mocks base method.
+func (m *MockNetwork) IsSynced() bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "IsSynced")
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// IsSynced indicates an expected call of IsSynced.
+func (mr *MockNetworkMockRecorder) IsSynced() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsSynced", reflect.TypeOf((*MockNetwork)(nil).IsSynced))
+}
+
+// ReportPeer mocks base method.
+func (m *MockNetwork) ReportPeer(arg0 peerset.ReputationChange, arg1 peer.ID) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "ReportPeer", arg0, arg1)
+}
+
+// ReportPeer indicates an expected call of ReportPeer.
+func (mr *MockNetworkMockRecorder) ReportPeer(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReportPeer", reflect.TypeOf((*MockNetwork)(nil).ReportPeer), arg0, arg1)
+}
diff --git a/dot/rpc/modules/system_integration_test.go b/dot/rpc/modules/system_integration_test.go
index 3d14ea1a1a..8151d522f3 100644
--- a/dot/rpc/modules/system_integration_test.go
+++ b/dot/rpc/modules/system_integration_test.go
@@ -21,7 +21,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/ChainSafe/gossamer/dot/core"
-	coremocks "github.com/ChainSafe/gossamer/dot/core/mocks"
 	"github.com/ChainSafe/gossamer/dot/network"
 	"github.com/ChainSafe/gossamer/dot/rpc/modules/mocks"
 	"github.com/ChainSafe/gossamer/dot/state"
@@ -351,6 +350,8 @@ func setupSystemModule(t *testing.T) *SystemModule {
 	return NewSystemModule(net, nil, core, chain.Storage, txQueue, nil, nil)
 }
 
+//go:generate mockgen -destination=mock_network_test.go -package $GOPACKAGE github.com/ChainSafe/gossamer/dot/core Network
+
 func newCoreService(t *testing.T, srvc *state.Service) *core.Service {
 	// setup service
 	tt := trie.NewEmptyTrie()
@@ -373,8 +374,12 @@ func newCoreService(t *testing.T, srvc *state.Service) *core.Service {
 		srvc = newTestStateService(t)
 	}
 
-	mocknet := new(coremocks.Network)
-	mocknet.On("GossipMessage", mock.AnythingOfType("network.NotificationsMessage"))
+	ctrl := gomock.NewController(t)
+
+	mocknet := NewMockNetwork(ctrl)
+	mocknet.EXPECT().GossipMessage(
+		gomock.AssignableToTypeOf(new(network.TransactionMessage))).
+		AnyTimes()
 
 	digestHandlerMock := NewMockDigestHandler(nil)
 
diff --git a/lib/babe/state.go b/lib/babe/state.go
index ae620ebe82..fa724276ab 100644
--- a/lib/babe/state.go
+++ b/lib/babe/state.go
@@ -84,6 +84,8 @@ type DigestHandler interface {
 	HandleDigests(*types.Header)
 }
 
+//go:generate mockery --name BlockImportHandler --structname BlockImportHandler --case underscore --keeptree
+
 // BlockImportHandler is the interface for the handler of new blocks
 type BlockImportHandler interface {
 	HandleBlockProduced(block *types.Block, state *rtstorage.TrieState) error
diff --git a/lib/grandpa/state.go b/lib/grandpa/state.go
index 7ba4087449..8e5a0ebd93 100644
--- a/lib/grandpa/state.go
+++ b/lib/grandpa/state.go
@@ -54,8 +54,6 @@ type GrandpaState interface { //nolint:revive
 	GetPrecommits(round, setID uint64) ([]SignedVote, error)
 }
 
-//go:generate mockery --name DigestHandler --structname DigestHandler --case underscore --keeptree
-
 // DigestHandler is the interface required by GRANDPA for the digest handler
 type DigestHandler interface { // TODO: use GrandpaState instead (#1871)
 	NextGrandpaAuthorityChange() uint64