diff --git a/Makefile b/Makefile
index bb35ec896dd..03c96756abc 100644
--- a/Makefile
+++ b/Makefile
@@ -216,6 +216,7 @@ generate-mocks: install-mock-generators
 	mockery --name 'Storage' --dir=module/executiondatasync/tracker --case=underscore --output="module/executiondatasync/tracker/mock" --outpkg="mocktracker"
 	mockery --name 'ScriptExecutor' --dir=module/execution --case=underscore --output="module/execution/mock" --outpkg="mock"
 	mockery --name 'StorageSnapshot' --dir=fvm/storage/snapshot --case=underscore --output="fvm/storage/snapshot/mock" --outpkg="mock"
+	mockery --name 'WebsocketConnection' --dir=engine/access/rest/websockets --case=underscore --output="engine/access/rest/websockets/mock" --outpkg="mock"
 
 	#temporarily make insecure/ a non-module to allow mockery to create mocks
 	mv insecure/go.mod insecure/go2.mod
diff --git a/engine/access/rest/websockets/controller.go b/engine/access/rest/websockets/controller.go
index 642a454271c..b999641df08 100644
--- a/engine/access/rest/websockets/controller.go
+++ b/engine/access/rest/websockets/controller.go
@@ -66,7 +66,7 @@ func (c *Controller) HandleConnection(ctx context.Context) {
 	g, gCtx := errgroup.WithContext(ctx)
 
 	g.Go(func() error {
-		return c.readMessagesFromClient(gCtx)
+		return c.readMessages(gCtx)
 	})
 
 	g.Go(func() error {
@@ -74,7 +74,7 @@ func (c *Controller) HandleConnection(ctx context.Context) {
 	})
 
 	g.Go(func() error {
-		return c.writeMessagesToClient(gCtx)
+		return c.writeMessages(gCtx)
 	})
 
 	if err = g.Wait(); err != nil {
@@ -111,12 +111,12 @@ func (c *Controller) configureKeepalive() error {
 	return nil
 }
 
-// writeMessagesToClient reads a messages from communication channel and passes them on to a client WebSocket connection.
+// writeMessages reads a messages from communication channel and passes them on to a client WebSocket connection.
 // The communication channel is filled by data providers. Besides, the response limit tracker is involved in
 // write message regulation
 //
 // No errors are expected during normal operation. All errors are considered benign.
-func (c *Controller) writeMessagesToClient(ctx context.Context) error {
+func (c *Controller) writeMessages(ctx context.Context) error {
 	for {
 		select {
 		case <-ctx.Done():
@@ -143,11 +143,11 @@ func (c *Controller) writeMessagesToClient(ctx context.Context) error {
 	}
 }
 
-// readMessagesFromClient continuously reads messages from a client WebSocket connection,
+// readMessages continuously reads messages from a client WebSocket connection,
 // processes each message, and handles actions based on the message type.
 //
 // No errors are expected during normal operation. All errors are considered benign.
-func (c *Controller) readMessagesFromClient(ctx context.Context) error {
+func (c *Controller) readMessages(ctx context.Context) error {
 	for {
 		select {
 		case <-ctx.Done():
@@ -163,10 +163,12 @@ func (c *Controller) readMessagesFromClient(ctx context.Context) error {
 
 			_, validatedMsg, err := c.parseAndValidateMessage(msg)
 			if err != nil {
+				//TODO: write error to error channel
 				return fmt.Errorf("failed to parse and validate client message: %w", err)
 			}
 
 			if err := c.handleAction(ctx, validatedMsg); err != nil {
+				//TODO: write error to error channel
 				return fmt.Errorf("failed to handle message action: %w", err)
 			}
 		}
@@ -241,8 +243,16 @@ func (c *Controller) handleSubscribe(ctx context.Context, msg models.SubscribeMe
 
 	c.dataProviders.Add(dp.ID(), dp)
 
-	//TODO: return OK response to client
-	c.communicationChannel <- msg
+	//TODO: return correct OK response to client
+	response := models.SubscribeMessageResponse{
+		BaseMessageResponse: models.BaseMessageResponse{
+			Success: true,
+		},
+		Topic: dp.Topic(),
+		ID:    dp.ID().String(),
+	}
+
+	c.communicationChannel <- response
 
 	go func() {
 		err := dp.Run()
diff --git a/engine/access/rest/websockets/controller_test.go b/engine/access/rest/websockets/controller_test.go
index afc14d9823f..58228a46b46 100644
--- a/engine/access/rest/websockets/controller_test.go
+++ b/engine/access/rest/websockets/controller_test.go
@@ -7,6 +7,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
+	streammock "github.com/onflow/flow-go/engine/access/state_stream/mock"
+
 	"github.com/google/uuid"
 	"github.com/gorilla/websocket"
 	"github.com/rs/zerolog"
@@ -18,6 +22,8 @@ import (
 	dpmock "github.com/onflow/flow-go/engine/access/rest/websockets/data_providers/mock"
 	connectionmock "github.com/onflow/flow-go/engine/access/rest/websockets/mock"
 	"github.com/onflow/flow-go/engine/access/rest/websockets/models"
+	"github.com/onflow/flow-go/engine/access/state_stream/backend"
+	"github.com/onflow/flow-go/model/flow"
 	"github.com/onflow/flow-go/utils/unittest"
 )
 
@@ -30,6 +36,9 @@ type ControllerSuite struct {
 
 	connection          *connectionmock.WebsocketConnection
 	dataProviderFactory *dpmock.DataProviderFactory
+
+	streamApi    *streammock.API
+	streamConfig backend.Config
 }
 
 func TestControllerSuite(t *testing.T) {
@@ -39,10 +48,232 @@ func TestControllerSuite(t *testing.T) {
 // SetupTest initializes the test suite with required dependencies.
 func (s *ControllerSuite) SetupTest() {
 	s.logger = unittest.Logger()
-	s.config = Config{}
+	s.config = NewDefaultWebsocketConfig()
 
 	s.connection = connectionmock.NewWebsocketConnection(s.T())
 	s.dataProviderFactory = dpmock.NewDataProviderFactory(s.T())
+
+	s.streamApi = streammock.NewAPI(s.T())
+	s.streamConfig = backend.Config{}
+}
+
+// TestSubscribeRequest tests the subscribe to topic flow.
+// We emulate a request message from a client, and a response message from a controller.
+func (s *ControllerSuite) TestSubscribeRequest() {
+	s.T().Run("Happy path", func(t *testing.T) {
+		conn, dataProviderFactory, dataProvider := newControllerMocks(t)
+		controller := NewWebSocketController(s.logger, s.config, conn, dataProviderFactory)
+
+		dataProvider.
+			On("Run").
+			Run(func(args mock.Arguments) {}).
+			Return(nil).
+			Once()
+
+		subscribeRequest := models.SubscribeMessageRequest{
+			BaseMessageRequest: models.BaseMessageRequest{Action: "subscribe"},
+			Topic:              dp.BlocksTopic,
+			Arguments:          nil,
+		}
+
+		// Simulate receiving the subscription request from the client
+		conn.
+			On("ReadJSON", mock.Anything).
+			Run(func(args mock.Arguments) {
+				requestMsg, ok := args.Get(0).(*json.RawMessage)
+				require.True(t, ok)
+				subscribeRequestMessage, err := json.Marshal(subscribeRequest)
+				require.NoError(t, err)
+				*requestMsg = subscribeRequestMessage
+			}).
+			Return(nil).
+			Once()
+
+		// Channel to signal the test flow completion
+		done := make(chan struct{}, 1)
+
+		// Simulate writing a successful subscription response back to the client
+		conn.
+			On("WriteJSON", mock.Anything).
+			Return(func(msg interface{}) error {
+				response, ok := msg.(models.SubscribeMessageResponse)
+				require.True(t, ok)
+				require.True(t, response.Success)
+				close(done) // Signal that response has been sent
+				return websocket.ErrCloseSent
+			}).Once()
+
+		// Simulate client closing connection after receiving the response
+		conn.
+			On("ReadJSON", mock.Anything).
+			Return(func(interface{}) error {
+				<-done
+				return websocket.ErrCloseSent
+			}).Once()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		controller.HandleConnection(ctx)
+	})
+}
+
+// TestSubscribeBlocks tests the functionality for streaming blocks to a subscriber.
+func (s *ControllerSuite) TestSubscribeBlocks() {
+	s.T().Run("Stream one block", func(t *testing.T) {
+		conn, dataProviderFactory, dataProvider := newControllerMocks(t)
+		controller := NewWebSocketController(s.logger, s.config, conn, dataProviderFactory)
+
+		// Simulate data provider write a block to the controller
+		expectedBlock := unittest.BlockFixture()
+		dataProvider.
+			On("Run", mock.Anything).
+			Run(func(args mock.Arguments) {
+				controller.communicationChannel <- expectedBlock
+			}).
+			Return(nil).
+			Once()
+
+		done := make(chan struct{}, 1)
+		s.expectSubscriptionRequest(conn, done)
+		s.expectSubscriptionResponse(conn, true)
+
+		// Expect a valid block to be passed to WriteJSON.
+		// If we got to this point, the controller executed all its logic properly
+		var actualBlock flow.Block
+		conn.
+			On("WriteJSON", mock.Anything).
+			Return(func(msg interface{}) error {
+				block, ok := msg.(flow.Block)
+				require.True(t, ok)
+				actualBlock = block
+
+				close(done)
+				return websocket.ErrCloseSent
+			}).Once()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		controller.HandleConnection(ctx)
+		require.Equal(t, expectedBlock, actualBlock)
+	})
+
+	s.T().Run("Stream many blocks", func(t *testing.T) {
+		conn, dataProviderFactory, dataProvider := newControllerMocks(t)
+		controller := NewWebSocketController(s.logger, s.config, conn, dataProviderFactory)
+
+		// Simulate data provider writes some blocks to the controller
+		expectedBlocks := unittest.BlockFixtures(100)
+		dataProvider.
+			On("Run", mock.Anything).
+			Run(func(args mock.Arguments) {
+				for _, block := range expectedBlocks {
+					controller.communicationChannel <- *block
+				}
+			}).
+			Return(nil).
+			Once()
+
+		done := make(chan struct{}, 1)
+		s.expectSubscriptionRequest(conn, done)
+		s.expectSubscriptionResponse(conn, true)
+
+		i := 0
+		actualBlocks := make([]*flow.Block, len(expectedBlocks))
+
+		// Expect valid blocks to be passed to WriteJSON.
+		// If we got to this point, the controller executed all its logic properly
+		conn.
+			On("WriteJSON", mock.Anything).
+			Return(func(msg interface{}) error {
+				block, ok := msg.(flow.Block)
+				require.True(t, ok)
+
+				actualBlocks[i] = &block
+				i += 1
+
+				if i == len(expectedBlocks) {
+					close(done)
+					return websocket.ErrCloseSent
+				}
+
+				return nil
+			}).
+			Times(len(expectedBlocks))
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		controller.HandleConnection(ctx)
+		require.Equal(t, expectedBlocks, actualBlocks)
+	})
+}
+
+// newControllerMocks initializes mock WebSocket connection, data provider, and data provider factory.
+// The mocked functions are expected to be called in a case when a test is expected to reach WriteJSON function.
+func newControllerMocks(t *testing.T) (*connectionmock.WebsocketConnection, *dpmock.DataProviderFactory, *dpmock.DataProvider) {
+	conn := connectionmock.NewWebsocketConnection(t)
+	conn.On("Close").Return(nil).Once()
+	conn.On("SetReadDeadline", mock.Anything).Return(nil).Once()
+	conn.On("SetWriteDeadline", mock.Anything).Return(nil)
+	conn.On("SetPongHandler", mock.AnythingOfType("func(string) error")).Return(nil).Once()
+
+	id := uuid.New()
+	topic := dp.BlocksTopic
+	dataProvider := dpmock.NewDataProvider(t)
+	dataProvider.On("ID").Return(id)
+	dataProvider.On("Close").Return(nil)
+	dataProvider.On("Topic").Return(topic)
+
+	factory := dpmock.NewDataProviderFactory(t)
+	factory.
+		On("NewDataProvider", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+		Return(dataProvider, nil).
+		Once()
+
+	return conn, factory, dataProvider
+}
+
+// expectSubscriptionRequest mocks the client's subscription request.
+func (s *ControllerSuite) expectSubscriptionRequest(conn *connectionmock.WebsocketConnection, done <-chan struct{}) {
+	requestMessage := models.SubscribeMessageRequest{
+		BaseMessageRequest: models.BaseMessageRequest{Action: "subscribe"},
+		Topic:              dp.BlocksTopic,
+	}
+
+	// The very first message from a client is a request to subscribe to some topic
+	conn.On("ReadJSON", mock.Anything).
+		Run(func(args mock.Arguments) {
+			reqMsg, ok := args.Get(0).(*json.RawMessage)
+			require.True(s.T(), ok)
+			msg, err := json.Marshal(requestMessage)
+			require.NoError(s.T(), err)
+			*reqMsg = msg
+		}).
+		Return(nil).
+		Once()
+
+	// In the default case, no further communication is expected from the client.
+	// We wait for the writer routine to signal completion, allowing us to close the connection gracefully
+	conn.
+		On("ReadJSON", mock.Anything).
+		Return(func(msg interface{}) error {
+			<-done
+			return websocket.ErrCloseSent
+		})
+}
+
+// expectSubscriptionResponse mocks the subscription response sent to the client.
+func (s *ControllerSuite) expectSubscriptionResponse(conn *connectionmock.WebsocketConnection, success bool) {
+	conn.On("WriteJSON", mock.Anything).
+		Run(func(args mock.Arguments) {
+			response, ok := args.Get(0).(models.SubscribeMessageResponse)
+			require.True(s.T(), ok)
+			require.Equal(s.T(), success, response.Success)
+		}).
+		Return(nil).
+		Once()
 }
 
 // TestConfigureKeepaliveConnection ensures that the WebSocket connection is configured correctly.
@@ -259,8 +490,9 @@ func (s *ControllerSuite) initializeController() *Controller {
 // mockDataProviderSetup is a helper which mocks a blocks data provider setup.
 func (s *ControllerSuite) mockBlockDataProviderSetup(id uuid.UUID) *dpmock.DataProvider {
 	dataProvider := dpmock.NewDataProvider(s.T())
-	dataProvider.On("ID").Return(id).Once()
+	dataProvider.On("ID").Return(id).Twice()
 	dataProvider.On("Close").Return(nil).Once()
+	dataProvider.On("Topic").Return(dp.BlocksTopic).Once()
 	s.dataProviderFactory.On("NewDataProvider", mock.Anything, dp.BlocksTopic, mock.Anything, mock.Anything).
 		Return(dataProvider, nil).Once()
 	dataProvider.On("Run").Return(nil).Once()
diff --git a/storage/mock/iter_item.go b/storage/mock/iter_item.go
new file mode 100644
index 00000000000..5d699511fb8
--- /dev/null
+++ b/storage/mock/iter_item.go
@@ -0,0 +1,82 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mock
+
+import mock "github.com/stretchr/testify/mock"
+
+// IterItem is an autogenerated mock type for the IterItem type
+type IterItem struct {
+	mock.Mock
+}
+
+// Key provides a mock function with given fields:
+func (_m *IterItem) Key() []byte {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Key")
+	}
+
+	var r0 []byte
+	if rf, ok := ret.Get(0).(func() []byte); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]byte)
+		}
+	}
+
+	return r0
+}
+
+// KeyCopy provides a mock function with given fields: dst
+func (_m *IterItem) KeyCopy(dst []byte) []byte {
+	ret := _m.Called(dst)
+
+	if len(ret) == 0 {
+		panic("no return value specified for KeyCopy")
+	}
+
+	var r0 []byte
+	if rf, ok := ret.Get(0).(func([]byte) []byte); ok {
+		r0 = rf(dst)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]byte)
+		}
+	}
+
+	return r0
+}
+
+// Value provides a mock function with given fields: _a0
+func (_m *IterItem) Value(_a0 func([]byte) error) error {
+	ret := _m.Called(_a0)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Value")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(func([]byte) error) error); ok {
+		r0 = rf(_a0)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// NewIterItem creates a new instance of IterItem. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewIterItem(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *IterItem {
+	mock := &IterItem{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/storage/mock/iterator.go b/storage/mock/iterator.go
new file mode 100644
index 00000000000..1b094ac15e1
--- /dev/null
+++ b/storage/mock/iterator.go
@@ -0,0 +1,106 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mock
+
+import (
+	storage "github.com/onflow/flow-go/storage"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Iterator is an autogenerated mock type for the Iterator type
+type Iterator struct {
+	mock.Mock
+}
+
+// Close provides a mock function with given fields:
+func (_m *Iterator) Close() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Close")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// First provides a mock function with given fields:
+func (_m *Iterator) First() bool {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for First")
+	}
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// IterItem provides a mock function with given fields:
+func (_m *Iterator) IterItem() storage.IterItem {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for IterItem")
+	}
+
+	var r0 storage.IterItem
+	if rf, ok := ret.Get(0).(func() storage.IterItem); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(storage.IterItem)
+		}
+	}
+
+	return r0
+}
+
+// Next provides a mock function with given fields:
+func (_m *Iterator) Next() {
+	_m.Called()
+}
+
+// Valid provides a mock function with given fields:
+func (_m *Iterator) Valid() bool {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Valid")
+	}
+
+	var r0 bool
+	if rf, ok := ret.Get(0).(func() bool); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(bool)
+	}
+
+	return r0
+}
+
+// NewIterator creates a new instance of Iterator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewIterator(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Iterator {
+	mock := &Iterator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/storage/mock/reader.go b/storage/mock/reader.go
new file mode 100644
index 00000000000..f9b15b532b5
--- /dev/null
+++ b/storage/mock/reader.go
@@ -0,0 +1,98 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mock
+
+import (
+	io "io"
+
+	storage "github.com/onflow/flow-go/storage"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Reader is an autogenerated mock type for the Reader type
+type Reader struct {
+	mock.Mock
+}
+
+// Get provides a mock function with given fields: key
+func (_m *Reader) Get(key []byte) ([]byte, io.Closer, error) {
+	ret := _m.Called(key)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Get")
+	}
+
+	var r0 []byte
+	var r1 io.Closer
+	var r2 error
+	if rf, ok := ret.Get(0).(func([]byte) ([]byte, io.Closer, error)); ok {
+		return rf(key)
+	}
+	if rf, ok := ret.Get(0).(func([]byte) []byte); ok {
+		r0 = rf(key)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]byte)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func([]byte) io.Closer); ok {
+		r1 = rf(key)
+	} else {
+		if ret.Get(1) != nil {
+			r1 = ret.Get(1).(io.Closer)
+		}
+	}
+
+	if rf, ok := ret.Get(2).(func([]byte) error); ok {
+		r2 = rf(key)
+	} else {
+		r2 = ret.Error(2)
+	}
+
+	return r0, r1, r2
+}
+
+// NewIter provides a mock function with given fields: startPrefix, endPrefix, ops
+func (_m *Reader) NewIter(startPrefix []byte, endPrefix []byte, ops storage.IteratorOption) (storage.Iterator, error) {
+	ret := _m.Called(startPrefix, endPrefix, ops)
+
+	if len(ret) == 0 {
+		panic("no return value specified for NewIter")
+	}
+
+	var r0 storage.Iterator
+	var r1 error
+	if rf, ok := ret.Get(0).(func([]byte, []byte, storage.IteratorOption) (storage.Iterator, error)); ok {
+		return rf(startPrefix, endPrefix, ops)
+	}
+	if rf, ok := ret.Get(0).(func([]byte, []byte, storage.IteratorOption) storage.Iterator); ok {
+		r0 = rf(startPrefix, endPrefix, ops)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(storage.Iterator)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func([]byte, []byte, storage.IteratorOption) error); ok {
+		r1 = rf(startPrefix, endPrefix, ops)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// NewReader creates a new instance of Reader. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewReader(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Reader {
+	mock := &Reader{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/storage/mock/reader_batch_writer.go b/storage/mock/reader_batch_writer.go
new file mode 100644
index 00000000000..c64c340704e
--- /dev/null
+++ b/storage/mock/reader_batch_writer.go
@@ -0,0 +1,72 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mock
+
+import (
+	storage "github.com/onflow/flow-go/storage"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// ReaderBatchWriter is an autogenerated mock type for the ReaderBatchWriter type
+type ReaderBatchWriter struct {
+	mock.Mock
+}
+
+// AddCallback provides a mock function with given fields: _a0
+func (_m *ReaderBatchWriter) AddCallback(_a0 func(error)) {
+	_m.Called(_a0)
+}
+
+// GlobalReader provides a mock function with given fields:
+func (_m *ReaderBatchWriter) GlobalReader() storage.Reader {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for GlobalReader")
+	}
+
+	var r0 storage.Reader
+	if rf, ok := ret.Get(0).(func() storage.Reader); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(storage.Reader)
+		}
+	}
+
+	return r0
+}
+
+// Writer provides a mock function with given fields:
+func (_m *ReaderBatchWriter) Writer() storage.Writer {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Writer")
+	}
+
+	var r0 storage.Writer
+	if rf, ok := ret.Get(0).(func() storage.Writer); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(storage.Writer)
+		}
+	}
+
+	return r0
+}
+
+// NewReaderBatchWriter creates a new instance of ReaderBatchWriter. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewReaderBatchWriter(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *ReaderBatchWriter {
+	mock := &ReaderBatchWriter{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
diff --git a/storage/mock/writer.go b/storage/mock/writer.go
new file mode 100644
index 00000000000..f80b206d39c
--- /dev/null
+++ b/storage/mock/writer.go
@@ -0,0 +1,81 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mock
+
+import (
+	storage "github.com/onflow/flow-go/storage"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Writer is an autogenerated mock type for the Writer type
+type Writer struct {
+	mock.Mock
+}
+
+// Delete provides a mock function with given fields: key
+func (_m *Writer) Delete(key []byte) error {
+	ret := _m.Called(key)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Delete")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func([]byte) error); ok {
+		r0 = rf(key)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// DeleteByRange provides a mock function with given fields: globalReader, startPrefix, endPrefix
+func (_m *Writer) DeleteByRange(globalReader storage.Reader, startPrefix []byte, endPrefix []byte) error {
+	ret := _m.Called(globalReader, startPrefix, endPrefix)
+
+	if len(ret) == 0 {
+		panic("no return value specified for DeleteByRange")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(storage.Reader, []byte, []byte) error); ok {
+		r0 = rf(globalReader, startPrefix, endPrefix)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Set provides a mock function with given fields: k, v
+func (_m *Writer) Set(k []byte, v []byte) error {
+	ret := _m.Called(k, v)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Set")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func([]byte, []byte) error); ok {
+		r0 = rf(k, v)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// NewWriter creates a new instance of Writer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewWriter(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Writer {
+	mock := &Writer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}