From 63f3cfe1b5093d5729d22d8af3d1eed8fffc703e Mon Sep 17 00:00:00 2001
From: Zhengke Zhou <madzhou1@gmail.com>
Date: Fri, 28 Mar 2025 16:45:34 +0800
Subject: [PATCH] sdk/trace: Manage trace environment variables in testing
 package (#6552)

Fix: #6542

Co-authored-by: Damien Mathieu <42@dmathieu.com>
---
 sdk/trace/batch_span_processor_test.go | 11 ----------
 sdk/trace/provider_test.go             | 29 ++++++++------------------
 sdk/trace/span_limits_test.go          |  4 ----
 sdk/trace/trace_test.go                |  9 ++------
 4 files changed, 11 insertions(+), 42 deletions(-)

diff --git a/sdk/trace/batch_span_processor_test.go b/sdk/trace/batch_span_processor_test.go
index 82cd1db0c..23e84138d 100644
--- a/sdk/trace/batch_span_processor_test.go
+++ b/sdk/trace/batch_span_processor_test.go
@@ -16,7 +16,6 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"go.opentelemetry.io/otel/sdk/internal/env"
-	ottest "go.opentelemetry.io/otel/sdk/internal/internaltest"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -247,16 +246,6 @@ func TestNewBatchSpanProcessorWithEnvOptions(t *testing.T) {
 		},
 	}
 
-	envStore := ottest.NewEnvStore()
-	envStore.Record(env.BatchSpanProcessorScheduleDelayKey)
-	envStore.Record(env.BatchSpanProcessorExportTimeoutKey)
-	envStore.Record(env.BatchSpanProcessorMaxQueueSizeKey)
-	envStore.Record(env.BatchSpanProcessorMaxExportBatchSizeKey)
-
-	defer func() {
-		require.NoError(t, envStore.Restore())
-	}()
-
 	for _, option := range options {
 		t.Run(option.name, func(t *testing.T) {
 			for k, v := range option.envs {
diff --git a/sdk/trace/provider_test.go b/sdk/trace/provider_test.go
index fbe63d319..577a5307d 100644
--- a/sdk/trace/provider_test.go
+++ b/sdk/trace/provider_test.go
@@ -14,10 +14,14 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"go.opentelemetry.io/otel/attribute"
-	ottest "go.opentelemetry.io/otel/sdk/internal/internaltest"
 	"go.opentelemetry.io/otel/trace"
 )
 
+const (
+	envTracesSampler    = "OTEL_TRACES_SAMPLER"
+	envTracesSamplerArg = "OTEL_TRACES_SAMPLER_ARG"
+)
+
 type basicSpanProcessor struct {
 	flushed             bool
 	closed              bool
@@ -313,19 +317,11 @@ func TestTracerProviderSamplerConfigFromEnv(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.sampler, func(t *testing.T) {
-			envVars := map[string]string{
-				"OTEL_TRACES_SAMPLER": test.sampler,
-			}
+			t.Setenv(envTracesSampler, test.sampler)
 
 			if test.samplerArg != "" {
-				envVars["OTEL_TRACES_SAMPLER_ARG"] = test.samplerArg
+				t.Setenv(envTracesSamplerArg, test.samplerArg)
 			}
-			envStore, err := ottest.SetEnvVariables(envVars)
-			require.NoError(t, err)
-			t.Cleanup(func() {
-				handler.Reset()
-				require.NoError(t, envStore.Restore())
-			})
 
 			stp := NewTracerProvider(WithSyncer(NewTestExporter()))
 			assert.Equal(t, test.description, stp.sampler.Description())
@@ -337,15 +333,8 @@ func TestTracerProviderSamplerConfigFromEnv(t *testing.T) {
 
 			if test.argOptional {
 				t.Run("invalid sampler arg", func(t *testing.T) {
-					envStore, err := ottest.SetEnvVariables(map[string]string{
-						"OTEL_TRACES_SAMPLER":     test.sampler,
-						"OTEL_TRACES_SAMPLER_ARG": "invalid-ignored-string",
-					})
-					require.NoError(t, err)
-					t.Cleanup(func() {
-						handler.Reset()
-						require.NoError(t, envStore.Restore())
-					})
+					t.Setenv(envTracesSampler, test.sampler)
+					t.Setenv(envTracesSamplerArg, "invalid-ignored-string")
 
 					stp := NewTracerProvider(WithSyncer(NewTestExporter()))
 					t.Cleanup(func() {
diff --git a/sdk/trace/span_limits_test.go b/sdk/trace/span_limits_test.go
index e6b90d052..871a4fe4b 100644
--- a/sdk/trace/span_limits_test.go
+++ b/sdk/trace/span_limits_test.go
@@ -12,7 +12,6 @@ import (
 
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/sdk/internal/env"
-	ottest "go.opentelemetry.io/otel/sdk/internal/internaltest"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -110,10 +109,7 @@ func TestSettingSpanLimits(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			if test.env != nil {
-				es := ottest.NewEnvStore()
-				t.Cleanup(func() { require.NoError(t, es.Restore()) })
 				for k, v := range test.env {
-					es.Record(k)
 					t.Setenv(k, v)
 				}
 			}
diff --git a/sdk/trace/trace_test.go b/sdk/trace/trace_test.go
index 994b6c2bf..2bf3dc33e 100644
--- a/sdk/trace/trace_test.go
+++ b/sdk/trace/trace_test.go
@@ -29,7 +29,7 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-const envVar = "OTEL_RESOURCE_ATTRIBUTES"
+const envVarResourceAttributes = "OTEL_RESOURCE_ATTRIBUTES"
 
 type storingHandler struct {
 	errs []error
@@ -1379,12 +1379,7 @@ func mergeResource(t *testing.T, r1, r2 *resource.Resource) *resource.Resource {
 }
 
 func TestWithResource(t *testing.T) {
-	store, err := ottest.SetEnvVariables(map[string]string{
-		envVar: "key=value,rk5=7",
-	})
-	require.NoError(t, err)
-	defer func() { require.NoError(t, store.Restore()) }()
-
+	t.Setenv(envVarResourceAttributes, "key=value,rk5=7")
 	cases := []struct {
 		name    string
 		options []TracerProviderOption