light: remove global context from tests (#7505)

light/client_test.go

@@ -27,7 +27,6 @@ const (
 )
 
 var (
-	ctx      = context.Background()
 	keys     = genPrivKeys(4)
 	vals     = keys.ToValidators(20, 10)
 	bTime, _ = time.Parse(time.RFC3339, "2006-01-02T15:04:05Z")
@@ -209,6 +208,9 @@ func TestClient_SequentialVerification(t *testing.T) {
 	for _, tc := range testCases {
 		testCase := tc
 		t.Run(testCase.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
 			logger := log.NewTestingLogger(t)
 
 			mockNode := mockNodeFromHeadersAndVals(testCase.otherHeaders, testCase.vals)
@@ -328,9 +330,14 @@ func TestClient_SkippingVerification(t *testing.T) {
 		},
 	}
 
+	bctx, bcancel := context.WithCancel(context.Background())
+	defer bcancel()
+
 	for _, tc := range testCases {
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(bctx)
+			defer cancel()
 			logger := log.NewTestingLogger(t)
 
 			mockNode := mockNodeFromHeadersAndVals(tc.otherHeaders, tc.vals)
@@ -382,6 +389,9 @@ func TestClientLargeBisectionVerification(t *testing.T) {
 
 	mockNode.On("LightBlock", mock.Anything, int64(0)).Return(lastBlock, nil)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	trustedLightBlock, err := mockNode.LightBlock(ctx, int64(200))
 	require.NoError(t, err)
 	c, err := light.NewClient(
@@ -410,6 +420,9 @@ func TestClientLargeBisectionVerification(t *testing.T) {
 }
 
 func TestClientBisectionBetweenTrustedHeaders(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mockFullNode := mockNodeFromHeadersAndVals(headerSet, valSet)
 	c, err := light.NewClient(
 		ctx,
@@ -440,6 +453,8 @@ func TestClientBisectionBetweenTrustedHeaders(t *testing.T) {
 }
 
 func TestClient_Cleanup(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	mockFullNode := &provider_mocks.Provider{}
@@ -469,8 +484,14 @@ func TestClient_Cleanup(t *testing.T) {
 
 // trustedHeader.Height == options.Height
 func TestClientRestoresTrustedHeaderAfterStartup(t *testing.T) {
+	bctx, bcancel := context.WithCancel(context.Background())
+	defer bcancel()
+
 	// 1. options.Hash == trustedHeader.Hash
 	t.Run("hashes should match", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(bctx)
+		defer cancel()
+
 		logger := log.NewTestingLogger(t)
 
 		mockNode := &provider_mocks.Provider{}
@@ -499,6 +520,9 @@ func TestClientRestoresTrustedHeaderAfterStartup(t *testing.T) {
 
 	// 2. options.Hash != trustedHeader.Hash
 	t.Run("hashes should not match", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(bctx)
+		defer cancel()
+
 		trustedStore := dbs.New(dbm.NewMemDB())
 		err := trustedStore.SaveLightBlock(l1)
 		require.NoError(t, err)
@@ -537,6 +561,9 @@ func TestClientRestoresTrustedHeaderAfterStartup(t *testing.T) {
 }
 
 func TestClient_Update(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mockFullNode := &provider_mocks.Provider{}
 	mockFullNode.On("LightBlock", mock.Anything, int64(0)).Return(l3, nil)
 	mockFullNode.On("LightBlock", mock.Anything, int64(1)).Return(l1, nil)
@@ -566,6 +593,8 @@ func TestClient_Update(t *testing.T) {
 }
 
 func TestClient_Concurrency(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	mockFullNode := &provider_mocks.Provider{}
@@ -613,6 +642,9 @@ func TestClient_Concurrency(t *testing.T) {
 }
 
 func TestClient_AddProviders(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mockFullNode := mockNodeFromHeadersAndVals(map[int64]*types.SignedHeader{
 		1: h1,
 		2: h2,
@@ -650,6 +682,9 @@ func TestClient_AddProviders(t *testing.T) {
 }
 
 func TestClientReplacesPrimaryWithWitnessIfPrimaryIsUnavailable(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mockFullNode := &provider_mocks.Provider{}
 	mockFullNode.On("LightBlock", mock.Anything, mock.Anything).Return(l1, nil)
 
@@ -683,6 +718,9 @@ func TestClientReplacesPrimaryWithWitnessIfPrimaryIsUnavailable(t *testing.T) {
 }
 
 func TestClientReplacesPrimaryWithWitnessIfPrimaryDoesntHaveBlock(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	mockFullNode := &provider_mocks.Provider{}
 	mockFullNode.On("LightBlock", mock.Anything, mock.Anything).Return(l1, nil)
 
@@ -711,6 +749,8 @@ func TestClientReplacesPrimaryWithWitnessIfPrimaryDoesntHaveBlock(t *testing.T)
 }
 
 func TestClient_BackwardsVerification(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	{
@@ -856,6 +896,9 @@ func TestClientRemovesWitnessIfItSendsUsIncorrectHeader(t *testing.T) {
 
 	mockFullNode := mockNodeFromHeadersAndVals(headerSet, valSet)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	lb1, _ := mockBadNode1.LightBlock(ctx, 2)
 	require.NotEqual(t, lb1.Hash(), l1.Hash())
 
@@ -917,6 +960,9 @@ func TestClient_TrustedValidatorSet(t *testing.T) {
 			2: vals,
 		})
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	c, err := light.NewClient(
 		ctx,
 		chainID,
@@ -949,6 +995,8 @@ func TestClientPrunesHeadersAndValidatorSets(t *testing.T) {
 			0: vals,
 		})
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	c, err := light.NewClient(
@@ -1036,6 +1084,9 @@ func TestClientEnsureValidHeadersAndValSets(t *testing.T) {
 	for i, tc := range testCases {
 		testCase := tc
 		t.Run(fmt.Sprintf("case: %d", i), func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
 			mockBadNode := mockNodeFromHeadersAndVals(testCase.headers, testCase.vals)
 			if testCase.errorToThrow != nil {
 				mockBadNode.On("LightBlock", mock.Anything, testCase.errorHeight).Return(nil, testCase.errorToThrow)

light/detector_test.go

@@ -2,6 +2,7 @@ package light_test
 
 import (
 	"bytes"
+	"context"
 	"testing"
 	"time"
 
@@ -31,6 +32,9 @@ func TestLightClientAttackEvidence_Lunatic(t *testing.T) {
 		primaryValidators = make(map[int64]*types.ValidatorSet, latestHeight)
 	)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	witnessHeaders, witnessValidators, chainKeys := genLightBlocksWithKeys(chainID, latestHeight, valSize, 2, bTime)
 
 	forgedKeys := chainKeys[divergenceHeight-1].ChangeKeys(3) // we change 3 out of the 5 validators (still 2/5 remain)
@@ -130,9 +134,15 @@ func TestLightClientAttackEvidence_Equivocation(t *testing.T) {
 		},
 	}
 
+	bctx, bcancel := context.WithCancel(context.Background())
+	defer bcancel()
+
 	for _, tc := range cases {
 		testCase := tc
 		t.Run(testCase.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(bctx)
+			defer cancel()
+
 			logger := log.NewTestingLogger(t)
 
 			// primary performs an equivocation attack
@@ -232,6 +242,8 @@ func TestLightClientAttackEvidence_ForwardLunatic(t *testing.T) {
 		primaryValidators = make(map[int64]*types.ValidatorSet, forgedHeight)
 	)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	witnessHeaders, witnessValidators, chainKeys := genLightBlocksWithKeys(chainID, latestHeight, valSize, 2, bTime)
@@ -380,6 +392,9 @@ func TestLightClientAttackEvidence_ForwardLunatic(t *testing.T) {
 // => light client returns an error upon creation because primary and witness
 // have a different view.
 func TestClientDivergentTraces1(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	headers, vals, _ := genLightBlocksWithKeys(chainID, 1, 5, 2, bTime)
 	mockPrimary := mockNodeFromHeadersAndVals(headers, vals)
 	firstBlock, err := mockPrimary.LightBlock(ctx, 1)
@@ -411,6 +426,8 @@ func TestClientDivergentTraces1(t *testing.T) {
 // 2. Two out of three nodes don't respond but the third has a header that matches
 // => verification should be successful and all the witnesses should remain
 func TestClientDivergentTraces2(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 	logger := log.NewTestingLogger(t)
 
 	headers, vals, _ := genLightBlocksWithKeys(chainID, 2, 5, 2, bTime)
@@ -451,6 +468,9 @@ func TestClientDivergentTraces3(t *testing.T) {
 	primaryHeaders, primaryVals, _ := genLightBlocksWithKeys(chainID, 2, 5, 2, bTime)
 	mockPrimary := mockNodeFromHeadersAndVals(primaryHeaders, primaryVals)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	firstBlock, err := mockPrimary.LightBlock(ctx, 1)
 	require.NoError(t, err)
 
@@ -491,6 +511,9 @@ func TestClientDivergentTraces4(t *testing.T) {
 	primaryHeaders, primaryVals, _ := genLightBlocksWithKeys(chainID, 2, 5, 2, bTime)
 	mockPrimary := mockNodeFromHeadersAndVals(primaryHeaders, primaryVals)
 
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
 	firstBlock, err := mockPrimary.LightBlock(ctx, 1)
 	require.NoError(t, err)