remove context.TODO() (#64)

Co-authored-by: Marko Baricevic <[email protected]>
rollkit · Feb 21, 2025 · bdae150 · bdae150
1 parent ae86f69
commit bdae150
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 8 deletions.
diff --git a/proxy/grpc/client_server_test.go b/proxy/grpc/client_server_test.go
@@ -71,7 +71,9 @@ func TestClientServer(t *testing.T) {
 		mockExec.On("InitChain", mock.Anything, expectedTime, initialHeight, chainID).
 			Return(stateRootHash, expectedMaxBytes, nil).Once()
 
-		stateRoot, maxBytes, err := client.InitChain(context.TODO(), genesisTime, initialHeight, chainID)
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		defer cancel()
+		stateRoot, maxBytes, err := client.InitChain(ctx, genesisTime, initialHeight, chainID)
 
 		require.NoError(t, err)
 		assert.Equal(t, stateRootHash, stateRoot)

diff --git a/proxy/grpc/proxy_test.go b/proxy/grpc/proxy_test.go
@@ -65,9 +65,12 @@ func (s *ProxyTestSuite) SetupTest() {
 	require.NoError(s.T(), err)
 
 	for i := 0; i < 10; i++ {
-		if _, err := client.GetTxs(context.TODO()); err == nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+		if _, err := client.GetTxs(ctx); err == nil {
+			cancel()
 			break
 		}
+		cancel()
 		time.Sleep(100 * time.Millisecond)
 	}
 

diff --git a/test/suite.go b/test/suite.go
@@ -28,7 +28,10 @@ func (s *ExecutorSuite) TestInitChain() {
 	initialHeight := uint64(1)
 	chainID := "test-chain"
 
-	stateRoot, maxBytes, err := s.Exec.InitChain(context.TODO(), genesisTime, initialHeight, chainID)
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	stateRoot, maxBytes, err := s.Exec.InitChain(ctx, genesisTime, initialHeight, chainID)
 	s.Require().NoError(err)
 	s.NotEqual(types.Hash{}, stateRoot)
 	s.Greater(maxBytes, uint64(0))
@@ -43,7 +46,11 @@ func (s *ExecutorSuite) TestGetTxs() {
 
 	s.TxInjector.InjectTx(tx1)
 	s.TxInjector.InjectTx(tx2)
-	txs, err := s.Exec.GetTxs(context.TODO())
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	txs, err := s.Exec.GetTxs(ctx)
 	s.Require().NoError(err)
 	s.Require().Len(txs, 2)
 	s.Require().Contains(txs, tx1)
@@ -63,21 +70,32 @@ func (s *ExecutorSuite) TestExecuteTxs() {
 	timestamp := time.Now().UTC()
 	prevStateRoot := types.Hash{1, 2, 3}
 
-	stateRoot, maxBytes, err := s.Exec.ExecuteTxs(context.TODO(), txs, blockHeight, timestamp, prevStateRoot)
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
+	stateRoot, maxBytes, err := s.Exec.ExecuteTxs(ctx, txs, blockHeight, timestamp, prevStateRoot)
 	s.Require().NoError(err)
 	s.NotEqual(types.Hash{}, stateRoot)
 	s.Greater(maxBytes, uint64(0))
 }
 
 // TestSetFinal tests SetFinal method.
 func (s *ExecutorSuite) TestSetFinal() {
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+
 	// finalizing invalid height must return error
-	err := s.Exec.SetFinal(context.TODO(), 1)
+	err := s.Exec.SetFinal(ctx, 1)
 	s.Require().Error(err)
 
-	_, _, err = s.Exec.ExecuteTxs(context.TODO(), nil, 2, time.Now(), types.Hash("test state"))
+	ctx2, cancel2 := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel2()
+	_, _, err = s.Exec.ExecuteTxs(ctx2, nil, 2, time.Now(), types.Hash("test state"))
 	s.Require().NoError(err)
-	err = s.Exec.SetFinal(context.TODO(), 2)
+
+	ctx3, cancel3 := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel3()
+	err = s.Exec.SetFinal(ctx3, 2)
 	s.Require().NoError(err)
 }