Skip to content

Commit 59a9550

Browse files
lucaspimentellucaspimentel
committed
Use [EnvironmentRestorer] in tests that set environment variables
Replace manual save/restore patterns with the [EnvironmentRestorer] attribute to prevent env var leaks between tests. Also allow the attribute on methods (not just classes) for more targeted usage. 🤖 Co-Authored-By: Claude Code <noreply@anthropic.com>
1 parent afb0834 commit 59a9550

5 files changed

Lines changed: 28 additions & 59 deletions

File tree

‎tracer/test/Datadog.Trace.TestHelpers/EnvironmentRestorerAttribute.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
namespace Datadog.Trace.TestHelpers;
1212

13-
[AttributeUsage(AttributeTargets.Class)]
13+
[AttributeUsage(AttributeTargets.Class | AttributeTargets.Method)]
1414
public class EnvironmentRestorerAttribute : BeforeAfterTestAttribute
1515
{
1616
private readonly string[] _variables;

‎tracer/test/Datadog.Trace.Tests/RuntimeMetrics/AzurePerformanceCountersListenerTests.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,15 @@
88
using System;
99
using Datadog.Trace.Configuration;
1010
using Datadog.Trace.RuntimeMetrics;
11+
using Datadog.Trace.TestHelpers;
1112
using Datadog.Trace.TestHelpers.Stats;
1213
using Datadog.Trace.Vendors.StatsdClient;
1314
using Moq;
1415
using Xunit;
1516

1617
namespace Datadog.Trace.Tests.RuntimeMetrics
1718
{
19+
[EnvironmentRestorer("WEBSITE_COUNTERS_CLR")]
1820
public class AzurePerformanceCountersListenerTests
1921
{
2022
[Fact]

‎tracer/test/Datadog.Trace.Tests/Telemetry/TelemetrySettingsAgentlessSettingsTests.cs‎

Lines changed: 3 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,38 +4,19 @@
44
// </copyright>
55

66
using System;
7-
using System.Collections.Generic;
87
using Datadog.Trace.Telemetry;
8+
using Datadog.Trace.TestHelpers;
99
using FluentAssertions;
1010
using Xunit;
1111

1212
namespace Datadog.Trace.Tests.Telemetry;
1313

1414
[Collection(nameof(EnvironmentVariablesTestCollection))]
15-
public class TelemetrySettingsAgentlessSettingsTests : IDisposable
15+
[EnvironmentRestorer("K_SERVICE", "CONTAINER_APP_NAME", "APPSVC_RUN_ZIP", "WEBSITE_APPSERVICEAPPLOGS_TRACE_ENABLED", "WEBSITE_SITE_NAME")]
16+
public class TelemetrySettingsAgentlessSettingsTests
1617
{
1718
private const string ApiKey = "some-key";
1819
private static readonly Uri Uri = new("http://localhost:8080");
19-
private static readonly string[] CloudVariables = { "K_SERVICE", "CONTAINER_APP_NAME", "APPSVC_RUN_ZIP", "WEBSITE_APPSERVICEAPPLOGS_TRACE_ENABLED", "WEBSITE_SITE_NAME" };
20-
private readonly Dictionary<string, string> _originalVariables = new();
21-
22-
public TelemetrySettingsAgentlessSettingsTests()
23-
{
24-
foreach (var variable in CloudVariables)
25-
{
26-
_originalVariables[variable] = Environment.GetEnvironmentVariable(variable);
27-
// clear variable
28-
Environment.SetEnvironmentVariable(variable, null);
29-
}
30-
}
31-
32-
public void Dispose()
33-
{
34-
foreach (var variable in _originalVariables)
35-
{
36-
Environment.SetEnvironmentVariable(variable.Key, variable.Value);
37-
}
38-
}
3920

4021
[Fact]
4122
public void CloudDetection_WhenNoCloudVariables()

‎tracer/test/Datadog.Trace.Tools.Runner.IntegrationTests/ConfigureCiCommandTests.cs‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,18 +23,17 @@ public class ConfigureCiCommandTests(ITestOutputHelper output)
2323
[InlineData("azp", @"##vso\[task.setvariable variable=(?<name>[A-Z1-9_]+);\](?<value>.*)")]
2424
[InlineData("jenkins", @"(?<name>[A-Z1-9_]+)=(?<value>.*)")]
2525
[InlineData("github", @"(?<name>[A-Z1-9_]+)=(?<value>.*)", "GITHUB_ENV")]
26+
[EnvironmentRestorer("GITHUB_ENV")]
2627
public void ConfigureCi(string ciProviderName, string pattern, string envKeyWithFilePath = null)
2728
{
2829
using var agent = MockTracerAgent.Create(output, TcpPortProvider.GetOpenPort());
2930
var agentUrl = $"http://localhost:{agent.Port}";
3031

3132
var commandLine = $"ci configure {ciProviderName} --dd-env TestEnv --dd-service TestService --dd-version TestVersion --tracer-home TestTracerHome --agent-url {agentUrl}";
3233

33-
string envKeyWithFilePathOriginalValue = null;
3434
string envKeyWithFilePathNewValue = null;
3535
if (!string.IsNullOrEmpty(envKeyWithFilePath))
3636
{
37-
envKeyWithFilePathOriginalValue = EnvironmentHelpers.GetEnvironmentVariable(envKeyWithFilePath);
3837
envKeyWithFilePathNewValue = Path.GetTempFileName();
3938
EnvironmentHelpers.SetEnvironmentVariable(envKeyWithFilePath, envKeyWithFilePathNewValue);
4039
}
@@ -71,11 +70,6 @@ public void ConfigureCi(string ciProviderName, string pattern, string envKeyWith
7170
environmentVariables.Should().Contain("DD_VERSION", "TestVersion");
7271
environmentVariables.Should().Contain("DD_DOTNET_TRACER_HOME", Path.GetFullPath("TestTracerHome"));
7372
environmentVariables.Should().Contain("DD_TRACE_AGENT_URL", agentUrl);
74-
75-
if (!string.IsNullOrEmpty(envKeyWithFilePath))
76-
{
77-
EnvironmentHelpers.SetEnvironmentVariable(envKeyWithFilePath, envKeyWithFilePathOriginalValue);
78-
}
7973
}
8074

8175
[SkippableTheory]

‎tracer/test/Datadog.Trace.Tools.Runner.IntegrationTests/LegacyCommandLineArgumentsTests.cs‎

Lines changed: 21 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -101,47 +101,39 @@ public void Run(bool withArguments)
101101
[Trait("RunOnWindows", "True")]
102102
[InlineData(' ')]
103103
[InlineData('=')]
104+
[EnvironmentRestorer("TF_BUILD")]
104105
public void SetCi(char separator)
105106
{
106-
var tfBuild = Environment.GetEnvironmentVariable("TF_BUILD");
107+
Environment.SetEnvironmentVariable("TF_BUILD", "1");
107108

108-
try
109-
{
110-
Environment.SetEnvironmentVariable("TF_BUILD", "1");
109+
using var console = ConsoleHelper.Redirect();
111110

112-
using var console = ConsoleHelper.Redirect();
111+
var commandLine = $"--set-ci --dd-env{separator}TestEnv --dd-service{separator}TestService --dd-version{separator}TestVersion --tracer-home{separator}TestTracerHome --agent-url{separator}TestAgentUrl --env-vars{separator}VAR1=A,VAR2=B";
113112

114-
var commandLine = $"--set-ci --dd-env{separator}TestEnv --dd-service{separator}TestService --dd-version{separator}TestVersion --tracer-home{separator}TestTracerHome --agent-url{separator}TestAgentUrl --env-vars{separator}VAR1=A,VAR2=B";
113+
var exitCode = Program.Main(commandLine.Split(' '));
115114

116-
var exitCode = Program.Main(commandLine.Split(' '));
115+
exitCode.Should().Be(0);
117116

118-
exitCode.Should().Be(0);
117+
var environmentVariables = new Dictionary<string, string>();
119118

120-
var environmentVariables = new Dictionary<string, string>();
119+
foreach (var line in console.ReadLines())
120+
{
121+
// ##vso[task.setvariable variable=DD_DOTNET_TRACER_HOME]TestTracerHome
122+
var match = Regex.Match(line, @"##vso\[task.setvariable variable=(?<name>[A-Z1-9_]+);\](?<value>.*)");
121123

122-
foreach (var line in console.ReadLines())
124+
if (match.Success)
123125
{
124-
// ##vso[task.setvariable variable=DD_DOTNET_TRACER_HOME]TestTracerHome
125-
var match = Regex.Match(line, @"##vso\[task.setvariable variable=(?<name>[A-Z1-9_]+);\](?<value>.*)");
126-
127-
if (match.Success)
128-
{
129-
environmentVariables.Add(match.Groups["name"].Value, match.Groups["value"].Value);
130-
}
126+
environmentVariables.Add(match.Groups["name"].Value, match.Groups["value"].Value);
131127
}
132-
133-
environmentVariables.Should().Contain("DD_ENV", "TestEnv");
134-
environmentVariables.Should().Contain("DD_SERVICE", "TestService");
135-
environmentVariables.Should().Contain("DD_VERSION", "TestVersion");
136-
environmentVariables.Should().Contain("DD_DOTNET_TRACER_HOME", Path.GetFullPath("TestTracerHome"));
137-
environmentVariables.Should().Contain("DD_TRACE_AGENT_URL", "TestAgentUrl");
138-
environmentVariables.Should().Contain("VAR1", "A");
139-
environmentVariables.Should().Contain("VAR2", "B");
140-
}
141-
finally
142-
{
143-
Environment.SetEnvironmentVariable("TF_BUILD", tfBuild);
144128
}
129+
130+
environmentVariables.Should().Contain("DD_ENV", "TestEnv");
131+
environmentVariables.Should().Contain("DD_SERVICE", "TestService");
132+
environmentVariables.Should().Contain("DD_VERSION", "TestVersion");
133+
environmentVariables.Should().Contain("DD_DOTNET_TRACER_HOME", Path.GetFullPath("TestTracerHome"));
134+
environmentVariables.Should().Contain("DD_TRACE_AGENT_URL", "TestAgentUrl");
135+
environmentVariables.Should().Contain("VAR1", "A");
136+
environmentVariables.Should().Contain("VAR2", "B");
145137
}
146138
}
147139
}

0 commit comments

Comments
 (0)