Skip to content

Commit c180092

Browse files
leodidoona-agent
andcommitted
refactor(test): replace os.Setenv with t.Setenv for automatic cleanup
Replace manual environment variable management with t.Setenv() which: - Automatically restores original values after test completion - Provides test isolation (each subtest gets clean environment) - Reduces boilerplate code by 92 lines - Eliminates manual defer cleanup logic This modernizes the test code to use Go 1.17+ testing features. Co-authored-by: Ona <no-reply@ona.com>
1 parent bd51715 commit c180092

File tree

1 file changed

+38
-130
lines changed

1 file changed

+38
-130
lines changed

pkg/leeway/signing/attestation_test.go

Lines changed: 38 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -630,30 +630,16 @@ func TestSignedAttestationResult_Structure(t *testing.T) {
630630

631631
// TestGetGitHubContext tests the environment variable extraction
632632
func TestGetGitHubContext(t *testing.T) {
633-
// Save original environment
634-
originalEnv := map[string]string{
635-
"GITHUB_RUN_ID": os.Getenv("GITHUB_RUN_ID"),
636-
"GITHUB_RUN_NUMBER": os.Getenv("GITHUB_RUN_NUMBER"),
637-
"GITHUB_ACTOR": os.Getenv("GITHUB_ACTOR"),
638-
"GITHUB_REPOSITORY": os.Getenv("GITHUB_REPOSITORY"),
639-
"GITHUB_REF": os.Getenv("GITHUB_REF"),
640-
"GITHUB_SHA": os.Getenv("GITHUB_SHA"),
641-
"GITHUB_SERVER_URL": os.Getenv("GITHUB_SERVER_URL"),
642-
"GITHUB_WORKFLOW_REF": os.Getenv("GITHUB_WORKFLOW_REF"),
643-
}
644-
645-
// Clean up after test
646-
defer func() {
647-
for k, v := range originalEnv {
648-
if v == "" {
649-
_ = os.Unsetenv(k)
650-
} else {
651-
_ = os.Setenv(k, v)
652-
}
653-
}
654-
}()
633+
// Set test environment (t.Setenv automatically handles cleanup)
634+
t.Setenv("GITHUB_RUN_ID", "test-run-id")
635+
t.Setenv("GITHUB_RUN_NUMBER", "test-run-number")
636+
t.Setenv("GITHUB_ACTOR", "test-actor")
637+
t.Setenv("GITHUB_REPOSITORY", "test-repo")
638+
t.Setenv("GITHUB_REF", "test-ref")
639+
t.Setenv("GITHUB_SHA", "test-sha")
640+
t.Setenv("GITHUB_SERVER_URL", "test-server")
641+
t.Setenv("GITHUB_WORKFLOW_REF", "test-workflow")
655642

656-
// Set test environment
657643
testEnv := map[string]string{
658644
"GITHUB_RUN_ID": "test-run-id",
659645
"GITHUB_RUN_NUMBER": "test-run-number",
@@ -665,10 +651,6 @@ func TestGetGitHubContext(t *testing.T) {
665651
"GITHUB_WORKFLOW_REF": "test-workflow",
666652
}
667653

668-
for k, v := range testEnv {
669-
_ = os.Setenv(k, v)
670-
}
671-
672654
// Test GetGitHubContext
673655
ctx := GetGitHubContext()
674656

@@ -684,39 +666,15 @@ func TestGetGitHubContext(t *testing.T) {
684666

685667
// TestGetGitHubContext_EmptyEnvironment tests with empty environment
686668
func TestGetGitHubContext_EmptyEnvironment(t *testing.T) {
687-
// Save original environment
688-
originalEnv := map[string]string{
689-
"GITHUB_RUN_ID": os.Getenv("GITHUB_RUN_ID"),
690-
"GITHUB_RUN_NUMBER": os.Getenv("GITHUB_RUN_NUMBER"),
691-
"GITHUB_ACTOR": os.Getenv("GITHUB_ACTOR"),
692-
"GITHUB_REPOSITORY": os.Getenv("GITHUB_REPOSITORY"),
693-
"GITHUB_REF": os.Getenv("GITHUB_REF"),
694-
"GITHUB_SHA": os.Getenv("GITHUB_SHA"),
695-
"GITHUB_SERVER_URL": os.Getenv("GITHUB_SERVER_URL"),
696-
"GITHUB_WORKFLOW_REF": os.Getenv("GITHUB_WORKFLOW_REF"),
697-
}
698-
699-
// Clean up after test
700-
defer func() {
701-
for k, v := range originalEnv {
702-
if v == "" {
703-
_ = os.Unsetenv(k)
704-
} else {
705-
_ = os.Setenv(k, v)
706-
}
707-
}
708-
}()
709-
710-
// Clear all GitHub environment variables
711-
githubVars := []string{
712-
"GITHUB_RUN_ID", "GITHUB_RUN_NUMBER", "GITHUB_ACTOR",
713-
"GITHUB_REPOSITORY", "GITHUB_REF", "GITHUB_SHA",
714-
"GITHUB_SERVER_URL", "GITHUB_WORKFLOW_REF",
715-
}
716-
717-
for _, v := range githubVars {
718-
_ = os.Unsetenv(v)
719-
}
669+
// Clear all GitHub environment variables (t.Setenv automatically handles cleanup)
670+
t.Setenv("GITHUB_RUN_ID", "")
671+
t.Setenv("GITHUB_RUN_NUMBER", "")
672+
t.Setenv("GITHUB_ACTOR", "")
673+
t.Setenv("GITHUB_REPOSITORY", "")
674+
t.Setenv("GITHUB_REF", "")
675+
t.Setenv("GITHUB_SHA", "")
676+
t.Setenv("GITHUB_SERVER_URL", "")
677+
t.Setenv("GITHUB_WORKFLOW_REF", "")
720678

721679
// Test GetGitHubContext with empty environment
722680
ctx := GetGitHubContext()
@@ -922,50 +880,32 @@ func (m *mockRemoteCache) HasFile(ctx context.Context, key string) (bool, error)
922880
// TestGetEnvOrDefault tests the environment variable helper
923881
// TestValidateSigstoreEnvironment tests Sigstore environment validation
924882
func TestValidateSigstoreEnvironment(t *testing.T) {
925-
// Save original environment
926-
originalEnv := map[string]string{
927-
"ACTIONS_ID_TOKEN_REQUEST_TOKEN": os.Getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN"),
928-
"ACTIONS_ID_TOKEN_REQUEST_URL": os.Getenv("ACTIONS_ID_TOKEN_REQUEST_URL"),
929-
"GITHUB_ACTIONS": os.Getenv("GITHUB_ACTIONS"),
930-
}
931-
932-
// Clean up after test
933-
defer func() {
934-
for k, v := range originalEnv {
935-
if v == "" {
936-
_ = os.Unsetenv(k)
937-
} else {
938-
_ = os.Setenv(k, v)
939-
}
940-
}
941-
}()
942-
943883
t.Run("missing required environment", func(t *testing.T) {
944-
// Clear all Sigstore environment variables
945-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
946-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_URL")
947-
_ = os.Unsetenv("GITHUB_ACTIONS")
884+
// Clear all Sigstore environment variables (t.Setenv automatically handles cleanup)
885+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "")
886+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "")
887+
t.Setenv("GITHUB_ACTIONS", "")
948888

949889
err := validateSigstoreEnvironment()
950890
assert.Error(t, err)
951891
assert.Contains(t, err.Error(), "ACTIONS_ID_TOKEN_REQUEST_TOKEN")
952892
})
953893

954894
t.Run("partial environment", func(t *testing.T) {
955-
// Set some but not all required variables
956-
_ = os.Setenv("GITHUB_ACTIONS", "true")
957-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
958-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_URL")
895+
// Set some but not all required variables (t.Setenv automatically handles cleanup)
896+
t.Setenv("GITHUB_ACTIONS", "true")
897+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "")
898+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "")
959899

960900
err := validateSigstoreEnvironment()
961901
assert.Error(t, err)
962902
})
963903

964904
t.Run("complete environment", func(t *testing.T) {
965-
// Set all required variables
966-
_ = os.Setenv("GITHUB_ACTIONS", "true")
967-
_ = os.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
968-
_ = os.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "https://test.url")
905+
// Set all required variables (t.Setenv automatically handles cleanup)
906+
t.Setenv("GITHUB_ACTIONS", "true")
907+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
908+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "https://test.url")
969909

970910
err := validateSigstoreEnvironment()
971911
assert.NoError(t, err)
@@ -1065,28 +1005,10 @@ func TestGenerateSignedSLSAAttestation_InvalidContext(t *testing.T) {
10651005

10661006
// TestSignProvenanceWithSigstore_EnvironmentValidation tests Sigstore environment validation
10671007
func TestSignProvenanceWithSigstore_EnvironmentValidation(t *testing.T) {
1068-
// Save original environment
1069-
originalEnv := map[string]string{
1070-
"ACTIONS_ID_TOKEN_REQUEST_TOKEN": os.Getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN"),
1071-
"ACTIONS_ID_TOKEN_REQUEST_URL": os.Getenv("ACTIONS_ID_TOKEN_REQUEST_URL"),
1072-
"GITHUB_ACTIONS": os.Getenv("GITHUB_ACTIONS"),
1073-
}
1074-
1075-
// Clean up after test
1076-
defer func() {
1077-
for k, v := range originalEnv {
1078-
if v == "" {
1079-
_ = os.Unsetenv(k)
1080-
} else {
1081-
_ = os.Setenv(k, v)
1082-
}
1083-
}
1084-
}()
1085-
1086-
// Clear Sigstore environment to trigger validation error
1087-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
1088-
_ = os.Unsetenv("ACTIONS_ID_TOKEN_REQUEST_URL")
1089-
_ = os.Unsetenv("GITHUB_ACTIONS")
1008+
// Clear Sigstore environment to trigger validation error (t.Setenv automatically handles cleanup)
1009+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "")
1010+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", "")
1011+
t.Setenv("GITHUB_ACTIONS", "")
10901012

10911013
artifactPath := createTestArtifact(t, "test content")
10921014
githubCtx := createMockGitHubContext()
@@ -1511,15 +1433,8 @@ func TestExtractBuilderIDFromOIDC(t *testing.T) {
15111433
server := tt.setupServer()
15121434
defer server.Close()
15131435

1514-
oldRequestURL := os.Getenv("ACTIONS_ID_TOKEN_REQUEST_URL")
1515-
oldRequestToken := os.Getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
1516-
defer func() {
1517-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", oldRequestURL)
1518-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", oldRequestToken)
1519-
}()
1520-
1521-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", server.URL)
1522-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
1436+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", server.URL)
1437+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
15231438

15241439
builderID, err := extractBuilderIDFromOIDC(context.Background(), tt.githubCtx)
15251440

@@ -1593,15 +1508,8 @@ func TestBuilderIDMatchesCertificateIdentity(t *testing.T) {
15931508
}))
15941509
defer server.Close()
15951510

1596-
oldRequestURL := os.Getenv("ACTIONS_ID_TOKEN_REQUEST_URL")
1597-
oldRequestToken := os.Getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN")
1598-
defer func() {
1599-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", oldRequestURL)
1600-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", oldRequestToken)
1601-
}()
1602-
1603-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", server.URL)
1604-
os.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
1511+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_URL", server.URL)
1512+
t.Setenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "test-token")
16051513

16061514
githubCtx := &GitHubContext{
16071515
ServerURL: "https://github.com",

0 commit comments

Comments
 (0)