fix type and unset usage

onflow · Dec 18, 2024 · 6810ae4 · 6810ae4
1 parent 89a59a7
commit 6810ae4
Showing 1 changed file with 66 additions and 29 deletions.
diff --git a/engine/access/rpc/connection/connection_test.go b/engine/access/rpc/connection/connection_test.go
@@ -92,7 +92,7 @@ func TestProxyExecutionAPI(t *testing.T) {
 	en.handler.
 		On("Ping",
 			testifymock.Anything,
-			testifymock.AnythingOfType("*access.PingRequest")).
+			testifymock.AnythingOfType("*execution.PingRequest")).
 		Return(expected, nil)
 
 	// create the factory
@@ -203,7 +203,7 @@ func TestProxyExecutionAPIConnectionReuse(t *testing.T) {
 	en.handler.
 		On("Ping",
 			testifymock.Anything,
-			testifymock.AnythingOfType("*access.PingRequest")).
+			testifymock.AnythingOfType("*execution.PingRequest")).
 		Return(expected, nil)
 
 	// create the factory
@@ -269,7 +269,7 @@ func TestExecutionNodeClientTimeout(t *testing.T) {
 	en.handler.
 		On("Ping",
 			testifymock.Anything,
-			testifymock.AnythingOfType("*access.PingRequest")).
+			testifymock.AnythingOfType("*execution.PingRequest")).
 		After(timeout+time.Second).
 		Return(resp, nil)
 
@@ -566,9 +566,14 @@ func TestExecutionNodeClientClosedGracefully(t *testing.T) {
 		req := &execution.PingRequest{}
 		resp := &execution.PingResponse{}
 		respSent := atomic.NewUint64(0)
-		en.handler.On("Ping", testifymock.Anything, req).Run(func(_ testifymock.Arguments) {
-			respSent.Inc()
-		}).Return(resp, nil)
+		en.handler.
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*execution.PingRequest")).
+			Run(func(_ testifymock.Arguments) {
+				respSent.Inc()
+			}).
+			Return(resp, nil)
 
 		// create the factory
 		connectionFactory := new(ConnectionFactoryImpl)
@@ -656,14 +661,18 @@ func TestEvictingCacheClients(t *testing.T) {
 	// Set up mock handlers for Ping and GetNetworkParameters
 	pingReq := &access.PingRequest{}
 	pingResp := &access.PingResponse{}
-	cn.handler.On("Ping", testifymock.Anything, pingReq).Return(
-		func(context.Context, *access.PingRequest) *access.PingResponse {
-			close(startPing)
-			<-returnFromPing // keeps request open until returnFromPing is closed
-			return pingResp
-		},
-		func(context.Context, *access.PingRequest) error { return nil },
-	)
+	cn.handler.
+		On("Ping",
+			testifymock.Anything,
+			testifymock.AnythingOfType("*access.PingRequest")).
+		Return(
+			func(context.Context, *access.PingRequest) *access.PingResponse {
+				close(startPing)
+				<-returnFromPing // keeps request open until returnFromPing is closed
+				return pingResp
+			},
+			func(context.Context, *access.PingRequest) error { return nil },
+		)
 
 	netReq := &access.GetNetworkParametersRequest{}
 	netResp := &access.GetNetworkParametersResponse{}
@@ -789,7 +798,9 @@ func TestConcurrentConnections(t *testing.T) {
 		requestCount := rapid.IntRange(50, 1000).Draw(tt, "r")
 		responsesSent := atomic.NewInt32(0)
 		en.handler.
-			On("Ping", testifymock.Anything, req).
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*execution.PingRequest")).
 			Return(func(_ context.Context, _ *execution.PingRequest) (*execution.PingResponse, error) {
 				time.Sleep(getSleep() * time.Microsecond)
 
@@ -932,7 +943,13 @@ func TestCircuitBreakerExecutionNode(t *testing.T) {
 		ctx := context.Background()
 
 		// Set up the handler mock to not respond within the requestTimeout.
-		en.handler.On("Ping", testifymock.Anything, req).After(2*requestTimeout).Return(resp, nil)
+		en.handler.
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*execution.PingRequest")).
+			After(2*requestTimeout).
+			Return(resp, nil).
+			Once()
 
 		// Call and measure the duration for the first invocation.
 		duration, err := callAndMeasurePingDuration(ctx)
@@ -941,28 +958,35 @@ func TestCircuitBreakerExecutionNode(t *testing.T) {
 
 		// Call and measure the duration for the second invocation (circuit breaker state is now "Open").
 		duration, err = callAndMeasurePingDuration(ctx)
-		assert.Equal(t, gobreaker.ErrOpenState, err)
+		assert.ErrorIs(t, err, gobreaker.ErrOpenState)
 		assert.Greater(t, requestTimeout, duration)
 
-		// Reset the mock Ping for the next invocation to return response without delay
-		en.handler.On("Ping", testifymock.Anything, req).Unset()
-		en.handler.On("Ping", testifymock.Anything, req).Return(resp, nil)
+		en.handler.
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*execution.PingRequest")).
+			Return(resp, nil).
+			Once()
 
 		// Wait until the circuit breaker transitions to the "HalfOpen" state.
 		time.Sleep(circuitBreakerRestoreTimeout + (500 * time.Millisecond))
 
 		// Call and measure the duration for the third invocation (circuit breaker state is now "HalfOpen").
 		duration, err = callAndMeasurePingDuration(ctx)
 		assert.Greater(t, requestTimeout, duration)
-		assert.Equal(t, nil, err)
+		assert.NoError(t, err)
 	})
 
 	for _, code := range successCodes {
 		t.Run(fmt.Sprintf("test error %s treated as a success for circuit breaker ", code.String()), func(t *testing.T) {
 			ctx := context.Background()
 
-			en.handler.On("Ping", testifymock.Anything, req).Unset()
-			en.handler.On("Ping", testifymock.Anything, req).Return(nil, status.Error(code, code.String()))
+			en.handler.
+				On("Ping",
+					testifymock.Anything,
+					testifymock.AnythingOfType("*execution.PingRequest")).
+				Return(nil, status.Error(code, code.String())).
+				Once()
 
 			duration, err := callAndMeasurePingDuration(ctx)
 			require.Error(t, err)
@@ -1038,7 +1062,13 @@ func TestCircuitBreakerCollectionNode(t *testing.T) {
 		ctx := context.Background()
 
 		// Set up the handler mock to not respond within the requestTimeout.
-		cn.handler.On("Ping", testifymock.Anything, req).After(2*requestTimeout).Return(resp, nil)
+		cn.handler.
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*access.PingRequest")).
+			After(2*requestTimeout).
+			Return(resp, nil).
+			Once()
 
 		// Call and measure the duration for the first invocation.
 		duration, err := callAndMeasurePingDuration(ctx)
@@ -1050,9 +1080,12 @@ func TestCircuitBreakerCollectionNode(t *testing.T) {
 		assert.Equal(t, gobreaker.ErrOpenState, err)
 		assert.Greater(t, requestTimeout, duration)
 
-		// Reset the mock Ping for the next invocation to return response without delay
-		cn.handler.On("Ping", testifymock.Anything, req).Unset()
-		cn.handler.On("Ping", testifymock.Anything, req).Return(resp, nil)
+		cn.handler.
+			On("Ping",
+				testifymock.Anything,
+				testifymock.AnythingOfType("*access.PingRequest")).
+			Return(resp, nil).
+			Once()
 
 		// Wait until the circuit breaker transitions to the "HalfOpen" state.
 		time.Sleep(circuitBreakerRestoreTimeout + (500 * time.Millisecond))
@@ -1067,8 +1100,12 @@ func TestCircuitBreakerCollectionNode(t *testing.T) {
 		t.Run(fmt.Sprintf("test error %s treated as a success for circuit breaker ", code.String()), func(t *testing.T) {
 			ctx := context.Background()
 
-			cn.handler.On("Ping", testifymock.Anything, req).Unset()
-			cn.handler.On("Ping", testifymock.Anything, req).Return(nil, status.Error(code, code.String()))
+			cn.handler.
+				On("Ping",
+					testifymock.Anything,
+					testifymock.AnythingOfType("*access.PingRequest")).
+				Return(nil, status.Error(code, code.String())).
+				Once()
 
 			duration, err := callAndMeasurePingDuration(ctx)
 			require.Error(t, err)