fix: use T.Context in tests (#2768)

julieqiu · web-flow · commit 7e7cd2ddc32b · 2025-11-04T11:51:56.000-05:00
T.Context returns a context that is canceled just before Cleanup-registered functions are called, and is intended to be used in tests. See https://pkg.go.dev/testing#T.Context.
diff --git a/internal/automation/cloudbuild_test.go b/internal/automation/cloudbuild_test.go
@@ -78,7 +78,7 @@ func TestRunCloudBuildTrigger(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx := t.Context()
 			client := &mockCloudBuildClient{
 				runError:      test.runError,
 				buildTriggers: make([]*cloudbuildpb.BuildTrigger, 0),
@@ -144,7 +144,7 @@ func TestFindTriggerIdByName(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx := t.Context()
 			client := &mockCloudBuildClient{
 				runError:      test.runError,
 				buildTriggers: test.buildTriggers,
@@ -209,7 +209,7 @@ func TestRunCloudBuildTriggerByName(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx := t.Context()
 			client := &mockCloudBuildClient{
 				runError:      test.runError,
 				buildTriggers: test.buildTriggers,
diff --git a/internal/automation/trigger_test.go b/internal/automation/trigger_test.go
@@ -184,7 +184,7 @@ func TestRunCommandWithClient(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx := t.Context()
 			client := &mockCloudBuildClient{
 				runError:      test.runError,
 				buildTriggers: test.buildTriggers,
@@ -334,7 +334,7 @@ func TestRunCommandWithConfig(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx := t.Context()
 			client := &mockCloudBuildClient{
 				runError:      test.runError,
 				buildTriggers: buildTriggers,
diff --git a/internal/container/java/execv/execv_test.go b/internal/container/java/execv/execv_test.go
@@ -15,7 +15,6 @@
 package execv
 
 import (
-	"context"
 	"errors"
 	"os/exec"
 	"strings"
@@ -56,7 +55,7 @@ func TestRun(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := Run(context.Background(), tt.args, ".")
+			err := Run(t.Context(), tt.args, ".")
 			if (err != nil) != tt.wantErr {
 				t.Fatalf("Run() error = %v, wantErr %v", err, tt.wantErr)
 			}
diff --git a/internal/container/java/generate/generator_test.go b/internal/container/java/generate/generator_test.go
@@ -274,7 +274,7 @@ java_gapic_library(
 			if err != nil {
 				t.Fatalf("failed to create generate config: %v", err)
 			}
-			if err := Generate(context.Background(), cfg); (err != nil) != tt.wantErr {
+			if err := Generate(t.Context(), cfg); (err != nil) != tt.wantErr {
 				t.Errorf("Generate() error = %v, wantErr %v", err, tt.wantErr)
 			}
 			if protocRunCount != tt.wantProtocRunCount {
diff --git a/internal/docker/docker_test.go b/internal/docker/docker_test.go
@@ -1020,7 +1020,7 @@ func TestReleaseStageRequestContent(t *testing.T) {
 		Output:          filepath.Join(tmpDir, "output"),
 		LibrarianConfig: &config.LibrarianConfig{},
 	}
-	if err := d.ReleaseStage(context.Background(), req); err != nil {
+	if err := d.ReleaseStage(t.Context(), req); err != nil {
 		t.Fatalf("d.ReleaseStage() failed: %v", err)
 	}
 }
diff --git a/internal/librarian/command_test.go b/internal/librarian/command_test.go
@@ -15,7 +15,6 @@
 package librarian
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -1564,7 +1563,7 @@ func TestCommitAndPush(t *testing.T) {
 				prBodyBuilder:     test.prBodyBuilder,
 			}
 
-			err := commitAndPush(context.Background(), commitInfo)
+			err := commitAndPush(t.Context(), commitInfo)
 
 			if test.wantErr {
 				if err == nil {
@@ -1746,7 +1745,7 @@ func TestAddLabelsToPullRequest(t *testing.T) {
 		t.Run(test.name, func(t *testing.T) {
 			client := test.mockGithubClient
 			prMetadata := test.prMetadata
-			err := addLabelsToPullRequest(context.Background(), client, test.wantPullRequestLabels, &prMetadata)
+			err := addLabelsToPullRequest(t.Context(), client, test.wantPullRequestLabels, &prMetadata)
 
 			if test.wantErr {
 				if err == nil {
diff --git a/internal/librarian/generate_command_test.go b/internal/librarian/generate_command_test.go
@@ -15,7 +15,6 @@
 package librarian
 
 import (
-	"context"
 	"errors"
 	"os"
 	"path/filepath"
@@ -358,7 +357,7 @@ func TestRunConfigureCommand(t *testing.T) {
 				}
 			}
 
-			_, err := r.runConfigureCommand(context.Background(), outputDir)
+			_, err := r.runConfigureCommand(t.Context(), outputDir)
 
 			if test.wantErr {
 				if err == nil {
@@ -949,7 +948,7 @@ func TestGenerateScenarios(t *testing.T) {
 				}
 			}
 
-			err := r.run(context.Background())
+			err := r.run(t.Context())
 			if test.wantErr {
 				if err == nil {
 					t.Fatalf("%s should return error", test.name)
@@ -1066,7 +1065,7 @@ func TestGenerateSingleLibraryCommand(t *testing.T) {
 				}
 			}
 
-			status, err := r.generateSingleLibrary(context.Background(), r.library, r.workRoot)
+			status, err := r.generateSingleLibrary(t.Context(), r.library, r.workRoot)
 			if test.wantErr {
 				if err == nil {
 					t.Fatalf("%s should return error", test.name)
diff --git a/internal/librarian/librarian_test.go b/internal/librarian/librarian_test.go
@@ -16,7 +16,6 @@ package librarian
 
 import (
 	"bytes"
-	"context"
 	"fmt"
 	"io"
 	"log/slog"
@@ -92,7 +91,7 @@ func TestVerboseFlag(t *testing.T) {
 			// Reset logger to default for test isolation.
 			slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, nil)))
 
-			_ = Run(context.Background(), test.args...)
+			_ = Run(t.Context(), test.args...)
 
 			// Restore stderr and read the output.
 			w.Close()
@@ -119,7 +118,7 @@ func TestVerboseFlag(t *testing.T) {
 }
 
 func TestGenerate_DefaultBehavior(t *testing.T) {
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// 1. Set up a mock repository with a state file
 	repo := newTestGitRepo(t)
diff --git a/internal/librarian/state_test.go b/internal/librarian/state_test.go
@@ -15,7 +15,6 @@
 package librarian
 
 import (
-	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -75,7 +74,7 @@ libraries:
       - src/b
     apis:
       - path: a/b/v1
-        service_config: 
+        service_config:
 `,
 			source:  "/non-existed-path",
 			want:    nil,
@@ -478,7 +477,7 @@ func TestLoadRepoStateFromGitHub(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			got, err := loadRepoStateFromGitHub(context.Background(), test.ghClient, test.branch)
+			got, err := loadRepoStateFromGitHub(t.Context(), test.ghClient, test.branch)
 			if test.wantErr {
 				if err == nil {
 					t.Fatal("loadRepoStateFromGitHub() should fail")
diff --git a/system_test.go b/system_test.go
@@ -15,7 +15,6 @@
 package librarian
 
 import (
-	"context"
 	"fmt"
 	"log/slog"
 	"os"
@@ -79,7 +78,7 @@ func TestGetRawContentSystem(t *testing.T) {
 					}
 
 					client := github.NewClient(testToken, repo)
-					got, err := client.GetRawContent(context.Background(), test.path, "main")
+					got, err := client.GetRawContent(t.Context(), test.path, "main")
 
 					if test.wantErr {
 						if err == nil {

Original file line number	Diff line number	Diff line change
`@@ -274,7 +274,7 @@ java_gapic_library(`
`274`	`274`	`if err != nil {`
`275`	`275`	`t.Fatalf("failed to create generate config: %v", err)`
`276`	`276`	`}`
`277`		`- if err := Generate(context.Background(), cfg); (err != nil) != tt.wantErr {`
	`277`	`+ if err := Generate(t.Context(), cfg); (err != nil) != tt.wantErr {`
`278`	`278`	`t.Errorf("Generate() error = %v, wantErr %v", err, tt.wantErr)`
`279`	`279`	`}`
`280`	`280`	`if protocRunCount != tt.wantProtocRunCount {`
Original file line number	Diff line number	Diff line change
`@@ -1020,7 +1020,7 @@ func TestReleaseStageRequestContent(t *testing.T) {`
`1020`	`1020`	`Output: filepath.Join(tmpDir, "output"),`
`1021`	`1021`	`LibrarianConfig: &config.LibrarianConfig{},`
`1022`	`1022`	`}`
`1023`		`- if err := d.ReleaseStage(context.Background(), req); err != nil {`
	`1023`	`+ if err := d.ReleaseStage(t.Context(), req); err != nil {`
`1024`	`1024`	`t.Fatalf("d.ReleaseStage() failed: %v", err)`
`1025`	`1025`	`}`
`1026`	`1026`	`}`
Original file line number	Diff line number	Diff line change
`@@ -15,7 +15,6 @@`
`15`	`15`	`package librarian`
`16`	`16`
`17`	`17`	`import (`
`18`		`- "context"`
`19`	`18`	`"errors"`
`20`	`19`	`"os"`
`21`	`20`	`"path/filepath"`
`@@ -358,7 +357,7 @@ func TestRunConfigureCommand(t *testing.T) {`
`358`	`357`	`}`
`359`	`358`	`}`
`360`	`359`
`361`		`- _, err := r.runConfigureCommand(context.Background(), outputDir)`
	`360`	`+ _, err := r.runConfigureCommand(t.Context(), outputDir)`
`362`	`361`
`363`	`362`	`if test.wantErr {`
`364`	`363`	`if err == nil {`
`@@ -949,7 +948,7 @@ func TestGenerateScenarios(t *testing.T) {`
`949`	`948`	`}`
`950`	`949`	`}`
`951`	`950`
`952`		`- err := r.run(context.Background())`
	`951`	`+ err := r.run(t.Context())`
`953`	`952`	`if test.wantErr {`
`954`	`953`	`if err == nil {`
`955`	`954`	`t.Fatalf("%s should return error", test.name)`
`@@ -1066,7 +1065,7 @@ func TestGenerateSingleLibraryCommand(t *testing.T) {`
`1066`	`1065`	`}`
`1067`	`1066`	`}`
`1068`	`1067`
`1069`		`- status, err := r.generateSingleLibrary(context.Background(), r.library, r.workRoot)`
	`1068`	`+ status, err := r.generateSingleLibrary(t.Context(), r.library, r.workRoot)`
`1070`	`1069`	`if test.wantErr {`
`1071`	`1070`	`if err == nil {`
`1072`	`1071`	`t.Fatalf("%s should return error", test.name)`