Merge pull request #6762 from The-K-R-O-K/illia-malachyn/6635-tests-for-websocket-controller

[Access] Add unit test for websocket controller

Author: Guitarheroua

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

engine/access/rest/websockets/controller.go

@@ -66,15 +66,15 @@ 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 {
 		return c.keepalive(gCtx)
 	})
 
 	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()

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()