Merge pull request #8012 from onflow/peter/speedup-websockets-keepalive-test

[Access] Speedup websockets keepalive test

Author: peterargue

engine/access/rest/websockets/config.go

@@ -64,3 +64,15 @@ func NewDefaultWebsocketConfig() Config {
 		InactivityTimeout:             DefaultInactivityTimeout,
 	}
 }
+
+type KeepaliveConfig struct {
+	PingPeriod time.Duration
+	PongWait   time.Duration
+}
+
+func DefaultKeepaliveConfig() KeepaliveConfig {
+	return KeepaliveConfig{
+		PingPeriod: PingPeriod,
+		PongWait:   PongWait,
+	}
+}

engine/access/rest/websockets/controller.go

@@ -136,6 +136,8 @@ type Controller struct {
 	dataProviderFactory dp.DataProviderFactory
 	dataProvidersGroup  *sync.WaitGroup
 	limiter             *rate.Limiter
+
+	keepaliveConfig KeepaliveConfig
 }
 
 func NewWebSocketController(
@@ -158,6 +160,7 @@ func NewWebSocketController(
 		dataProviderFactory: dataProviderFactory,
 		dataProvidersGroup:  &sync.WaitGroup{},
 		limiter:             limiter,
+		keepaliveConfig:     DefaultKeepaliveConfig(),
 	}
 }
 
@@ -213,13 +216,13 @@ func (c *Controller) configureKeepalive() error {
 	// The Pong handler itself only resets the read deadline after receiving a Pong.
 	// It doesn't set an initial deadline. The initial read deadline is crucial to prevent the server from waiting
 	// forever if the client doesn't send Pongs.
-	if err := c.conn.SetReadDeadline(time.Now().Add(PongWait)); err != nil {
+	if err := c.conn.SetReadDeadline(time.Now().Add(c.keepaliveConfig.PongWait)); err != nil {
 		return fmt.Errorf("failed to set the initial read deadline: %w", err)
 	}
 
 	// Establish a Pong handler which sets the handler for pong messages received from the peer.
 	c.conn.SetPongHandler(func(string) error {
-		return c.conn.SetReadDeadline(time.Now().Add(PongWait))
+		return c.conn.SetReadDeadline(time.Now().Add(c.keepaliveConfig.PongWait))
 	})
 
 	return nil
@@ -235,7 +238,7 @@ func (c *Controller) keepalive(ctx context.Context) error {
 		}
 	}()
 
-	pingTicker := time.NewTicker(PingPeriod)
+	pingTicker := time.NewTicker(c.keepaliveConfig.PingPeriod)
 	defer pingTicker.Stop()
 
 	for {

engine/access/rest/websockets/controller_test.go

@@ -894,6 +894,11 @@ func (s *WsControllerSuite) TestControllerShutdown() {
 }
 
 func (s *WsControllerSuite) TestKeepaliveRoutine() {
+	keepaliveConfig := KeepaliveConfig{
+		PingPeriod: time.Microsecond,
+		PongWait:   2 * time.Microsecond,
+	}
+
 	s.T().Run("Successfully pings connection n times", func(t *testing.T) {
 		conn := connmock.NewWebsocketConnection(t)
 		conn.On("Close").Return(nil).Once()
@@ -923,9 +928,9 @@ func (s *WsControllerSuite) TestKeepaliveRoutine() {
 
 		factory := dpmock.NewDataProviderFactory(t)
 		controller := NewWebSocketController(s.logger, s.wsConfig, conn, factory)
-		controller.HandleConnection(context.Background())
+		controller.keepaliveConfig = keepaliveConfig
 
-		conn.AssertExpectations(t)
+		controller.HandleConnection(context.Background())
 	})
 
 	s.T().Run("Error on write to closed connection", func(t *testing.T) {
@@ -938,15 +943,14 @@ func (s *WsControllerSuite) TestKeepaliveRoutine() {
 
 		factory := dpmock.NewDataProviderFactory(t)
 		controller := NewWebSocketController(s.logger, s.wsConfig, conn, factory)
+		controller.keepaliveConfig = keepaliveConfig
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 
 		err := controller.keepalive(ctx)
 		s.Require().Error(err)
 		s.Require().ErrorIs(expectedError, err)
-
-		conn.AssertExpectations(t)
 	})
 
 	s.T().Run("Error on write to open connection", func(t *testing.T) {
@@ -958,30 +962,28 @@ func (s *WsControllerSuite) TestKeepaliveRoutine() {
 
 		factory := dpmock.NewDataProviderFactory(t)
 		controller := NewWebSocketController(s.logger, s.wsConfig, conn, factory)
+		controller.keepaliveConfig = keepaliveConfig
 
 		ctx, cancel := context.WithCancel(context.Background())
 		defer cancel()
 
 		err := controller.keepalive(ctx)
 		s.Require().Error(err)
 		s.Require().ErrorContains(err, "error sending ping")
-
-		conn.AssertExpectations(t)
 	})
 
 	s.T().Run("Context cancelled", func(t *testing.T) {
 		conn := connmock.NewWebsocketConnection(t)
 		factory := dpmock.NewDataProviderFactory(t)
 		controller := NewWebSocketController(s.logger, s.wsConfig, conn, factory)
+		controller.keepaliveConfig = keepaliveConfig
 
 		ctx, cancel := context.WithCancel(context.Background())
 		cancel() // Immediately cancel the context
 
 		// Start the keepalive process with the context canceled
 		err := controller.keepalive(ctx)
 		s.Require().NoError(err)
-
-		conn.AssertExpectations(t) // Should not invoke WriteMessage after context cancellation
 	})
 }