From 742eb8f424832e68112e3f0ecf2db8e4c7a087a7 Mon Sep 17 00:00:00 2001
From: Jordan Krage <jmank88@gmail.com>
Date: Fri, 3 Jan 2025 14:06:59 -0600
Subject: [PATCH] core/chains/evm/client: fix test race (#15702)

---
 GNUmakefile                               |  2 +-
 core/chains/evm/client/rpc_client_test.go | 15 ++++++++-------
 core/internal/testutils/testutils.go      | 14 +-------------
 3 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/GNUmakefile b/GNUmakefile
index 08324b1fec4..bb200cc2cb9 100644
--- a/GNUmakefile
+++ b/GNUmakefile
@@ -183,7 +183,7 @@ modgraph:
 
 .PHONY: test-short
 test-short: ## Run 'go test -short' and suppress uninteresting output
-	go test -short ./... | grep -v "no test files" | grep -v "\(cached\)"
+	go test -short ./... | grep -v "\[no test files\]" | grep -v "\(cached\)"
 
 help:
 	@echo ""
diff --git a/core/chains/evm/client/rpc_client_test.go b/core/chains/evm/client/rpc_client_test.go
index 109a49d6e2f..d5286e9acf0 100644
--- a/core/chains/evm/client/rpc_client_test.go
+++ b/core/chains/evm/client/rpc_client_test.go
@@ -8,6 +8,7 @@ import (
 	"math/big"
 	"net/url"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -400,7 +401,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 	lggr := logger.Test(t)
 
 	type rpcServer struct {
-		Head *evmtypes.Head
+		Head atomic.Pointer[evmtypes.Head]
 		URL  *url.URL
 	}
 	createRPCServer := func() *rpcServer {
@@ -408,7 +409,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 		server.URL = testutils.NewWSServer(t, chainId, func(method string, params gjson.Result) (resp testutils.JSONRPCResponse) {
 			assert.Equal(t, "eth_getBlockByNumber", method)
 			if assert.True(t, params.IsArray()) && assert.Equal(t, "finalized", params.Array()[0].String()) {
-				head := server.Head
+				head := server.Head.Load()
 				jsonHead, err := json.Marshal(head)
 				if err != nil {
 					panic(fmt.Errorf("failed to marshal head: %w", err))
@@ -426,7 +427,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 	rpc := client.NewRPCClient(nodePoolCfg, lggr, server.URL, nil, "rpc", 1, chainId, commonclient.Primary, commonclient.QueryTimeout, commonclient.QueryTimeout, "")
 	require.NoError(t, rpc.Dial(ctx))
 	defer rpc.Close()
-	server.Head = &evmtypes.Head{Number: 128}
+	server.Head.Store(&evmtypes.Head{Number: 128})
 	// updates chain info
 	_, err := rpc.LatestFinalizedBlock(ctx)
 	require.NoError(t, err)
@@ -439,7 +440,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 	assert.Equal(t, int64(128), latest.FinalizedBlockNumber)
 
 	// lower block number does not update highestUserObservations
-	server.Head = &evmtypes.Head{Number: 127}
+	server.Head.Store(&evmtypes.Head{Number: 127})
 	_, err = rpc.LatestFinalizedBlock(ctx)
 	require.NoError(t, err)
 	latest, highestUserObservations = rpc.GetInterceptedChainInfo()
@@ -451,7 +452,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 	assert.Equal(t, int64(127), latest.FinalizedBlockNumber)
 
 	// health check flg prevents change in highestUserObservations
-	server.Head = &evmtypes.Head{Number: 256}
+	server.Head.Store(&evmtypes.Head{Number: 256})
 	_, err = rpc.LatestFinalizedBlock(commonclient.CtxAddHealthCheckFlag(ctx))
 	require.NoError(t, err)
 	latest, highestUserObservations = rpc.GetInterceptedChainInfo()
@@ -463,7 +464,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 	assert.Equal(t, int64(256), latest.FinalizedBlockNumber)
 
 	// subscription updates chain info
-	server.Head = &evmtypes.Head{Number: 512}
+	server.Head.Store(&evmtypes.Head{Number: 512})
 	ch, sub, err := rpc.SubscribeToFinalizedHeads(ctx)
 	require.NoError(t, err)
 	defer sub.Unsubscribe()
@@ -479,7 +480,7 @@ func TestRPCClient_LatestFinalizedBlock(t *testing.T) {
 
 	// health check subscription only updates latest
 	sub.Unsubscribe() // close previous one
-	server.Head = &evmtypes.Head{Number: 1024}
+	server.Head.Store(&evmtypes.Head{Number: 1024})
 	ch, sub, err = rpc.SubscribeToFinalizedHeads(commonclient.CtxAddHealthCheckFlag(ctx))
 	require.NoError(t, err)
 	defer sub.Unsubscribe()
diff --git a/core/internal/testutils/testutils.go b/core/internal/testutils/testutils.go
index 53b555c0e8b..c2f7351186e 100644
--- a/core/internal/testutils/testutils.go
+++ b/core/internal/testutils/testutils.go
@@ -126,19 +126,7 @@ func WaitTimeout(t *testing.T) time.Duration {
 
 // Context returns a context with the test's deadline, if available.
 func Context(tb testing.TB) context.Context {
-	ctx := context.Background()
-	var cancel func()
-	switch t := tb.(type) {
-	case *testing.T:
-		if d, ok := t.Deadline(); ok {
-			ctx, cancel = context.WithDeadline(ctx, d)
-		}
-	}
-	if cancel == nil {
-		ctx, cancel = context.WithCancel(ctx)
-	}
-	tb.Cleanup(cancel)
-	return ctx
+	return tests.Context(tb)
 }
 
 // MustParseURL parses the URL or fails the test