Skip to content

Commit 8b6d925

Browse files
hjgracahjgraca
authored
Merge pull request #835 from hjgraca/fix/override-lambda-console
chore: Fix Console output when running on Top Level Statements
2 parents 6d5a12c + af3419e commit 8b6d925

File tree

9 files changed

+296
-46
lines changed

9 files changed

+296
-46
lines changed

‎libraries/src/AWS.Lambda.Powertools.Common/Core/ISystemWrapper.cs

-11
Original file line numberDiff line numberDiff line change
@@ -59,15 +59,4 @@ public interface ISystemWrapper
5959
/// </summary>
6060
/// <param name="type"></param>
6161
void SetExecutionEnvironment<T>(T type);
62-
63-
/// <summary>
64-
/// Sets console output
65-
/// Useful for testing and checking the console output
66-
/// <code>
67-
/// var consoleOut = new StringWriter();
68-
/// SystemWrapper.Instance.SetOut(consoleOut);
69-
/// </code>
70-
/// </summary>
71-
/// <param name="writeTo">The TextWriter instance where to write to</param>
72-
void SetOut(TextWriter writeTo);
7362
}

‎libraries/src/AWS.Lambda.Powertools.Common/Core/SystemWrapper.cs

+68-11
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@ namespace AWS.Lambda.Powertools.Common;
2727
public class SystemWrapper : ISystemWrapper
2828
{
2929
private static IPowertoolsEnvironment _powertoolsEnvironment;
30+
private static bool _inTestMode = false;
31+
private static TextWriter _testOutputStream;
32+
private static bool _outputResetPerformed = false;
3033

3134
/// <summary>
3235
/// The instance
@@ -41,13 +44,11 @@ public SystemWrapper(IPowertoolsEnvironment powertoolsEnvironment)
4144
_powertoolsEnvironment = powertoolsEnvironment;
4245
_instance ??= this;
4346

44-
// Clear AWS SDK Console injected parameters StdOut and StdErr
45-
var standardOutput = new StreamWriter(Console.OpenStandardOutput());
46-
standardOutput.AutoFlush = true;
47-
Console.SetOut(standardOutput);
48-
var errordOutput = new StreamWriter(Console.OpenStandardError());
49-
errordOutput.AutoFlush = true;
50-
Console.SetError(errordOutput);
47+
if (!_inTestMode)
48+
{
49+
// Clear AWS SDK Console injected parameters in production only
50+
ResetConsoleOutput();
51+
}
5152
}
5253

5354
/// <summary>
@@ -72,7 +73,15 @@ public string GetEnvironmentVariable(string variable)
7273
/// <param name="value">The value.</param>
7374
public void Log(string value)
7475
{
75-
Console.Write(value);
76+
if (_inTestMode && _testOutputStream != null)
77+
{
78+
_testOutputStream.Write(value);
79+
}
80+
else
81+
{
82+
EnsureConsoleOutputOnce();
83+
Console.Write(value);
84+
}
7685
}
7786

7887
/// <summary>
@@ -81,7 +90,15 @@ public void Log(string value)
8190
/// <param name="value">The value.</param>
8291
public void LogLine(string value)
8392
{
84-
Console.WriteLine(value);
93+
if (_inTestMode && _testOutputStream != null)
94+
{
95+
_testOutputStream.WriteLine(value);
96+
}
97+
else
98+
{
99+
EnsureConsoleOutputOnce();
100+
Console.WriteLine(value);
101+
}
85102
}
86103

87104
/// <summary>
@@ -126,9 +143,20 @@ public void SetExecutionEnvironment<T>(T type)
126143
SetEnvironmentVariable(envName, envValue.ToString());
127144
}
128145

129-
/// <inheritdoc />
130-
public void SetOut(TextWriter writeTo)
146+
/// <summary>
147+
/// Sets console output
148+
/// Useful for testing and checking the console output
149+
/// <code>
150+
/// var consoleOut = new StringWriter();
151+
/// SystemWrapper.Instance.SetOut(consoleOut);
152+
/// </code>
153+
/// </summary>
154+
/// <param name="writeTo">The TextWriter instance where to write to</param>
155+
156+
public static void SetOut(TextWriter writeTo)
131157
{
158+
_testOutputStream = writeTo;
159+
_inTestMode = true;
132160
Console.SetOut(writeTo);
133161
}
134162

@@ -152,4 +180,33 @@ private string ParseAssemblyName(string assemblyName)
152180

153181
return $"{Constants.FeatureContextIdentifier}/{assemblyName}";
154182
}
183+
184+
private static void EnsureConsoleOutputOnce()
185+
{
186+
if (_outputResetPerformed) return;
187+
ResetConsoleOutput();
188+
_outputResetPerformed = true;
189+
}
190+
191+
private static void ResetConsoleOutput()
192+
{
193+
var standardOutput = new StreamWriter(Console.OpenStandardOutput());
194+
standardOutput.AutoFlush = true;
195+
Console.SetOut(standardOutput);
196+
var errorOutput = new StreamWriter(Console.OpenStandardError());
197+
errorOutput.AutoFlush = true;
198+
Console.SetError(errorOutput);
199+
}
200+
201+
public static void ClearOutputResetFlag()
202+
{
203+
_outputResetPerformed = false;
204+
}
205+
206+
// For test cleanup
207+
internal static void ResetTestMode()
208+
{
209+
_inTestMode = false;
210+
_testOutputStream = null;
211+
}
155212
}

‎libraries/tests/AWS.Lambda.Powertools.Common.Tests/Core/SystemWrapperTests.cs

+204
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
using System;
2+
using System.IO;
3+
using System.Reflection;
4+
using NSubstitute;
5+
using Xunit;
6+
7+
namespace AWS.Lambda.Powertools.Common.Tests;
8+
9+
[Collection("Sequential")]
10+
public class SystemWrapperTests : IDisposable
11+
{
12+
private readonly IPowertoolsEnvironment _mockEnvironment;
13+
private readonly StringWriter _testWriter;
14+
private readonly FieldInfo _outputResetPerformedField;
15+
16+
17+
public SystemWrapperTests()
18+
{
19+
_mockEnvironment = Substitute.For<IPowertoolsEnvironment>();
20+
_testWriter = new StringWriter();
21+
22+
// Get access to private field for testing
23+
_outputResetPerformedField = typeof(SystemWrapper).GetField("_outputResetPerformed",
24+
BindingFlags.NonPublic | BindingFlags.Static);
25+
26+
// Reset static state between tests
27+
SystemWrapper.ResetTestMode();
28+
_outputResetPerformedField.SetValue(null, false);
29+
}
30+
31+
[Fact]
32+
public void Log_InProductionMode_ResetsOutputOnce()
33+
{
34+
// Arrange
35+
var wrapper = new SystemWrapper(_mockEnvironment);
36+
var message1 = "First message";
37+
var message2 = "Second message";
38+
_outputResetPerformedField.SetValue(null, false);
39+
40+
// Act
41+
wrapper.Log(message1);
42+
bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
43+
wrapper.Log(message2);
44+
bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
45+
46+
// Assert
47+
Assert.True(afterFirstLog, "Flag should be set after first log");
48+
Assert.True(afterSecondLog, "Flag should remain set after second log");
49+
}
50+
51+
[Fact]
52+
public void LogLine_InProductionMode_ResetsOutputOnce()
53+
{
54+
// Arrange
55+
var wrapper = new SystemWrapper(_mockEnvironment);
56+
var message1 = "First line";
57+
var message2 = "Second line";
58+
_outputResetPerformedField.SetValue(null, false);
59+
60+
// Act
61+
wrapper.LogLine(message1);
62+
bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
63+
wrapper.LogLine(message2);
64+
bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
65+
66+
// Assert
67+
Assert.True(afterFirstLog, "Flag should be set after first LogLine");
68+
Assert.True(afterSecondLog, "Flag should remain set after second LogLine");
69+
}
70+
71+
[Fact]
72+
public void ClearOutputResetFlag_ResetsFlag_AllowsSubsequentReset()
73+
{
74+
// Arrange
75+
var wrapper = new SystemWrapper(_mockEnvironment);
76+
_outputResetPerformedField.SetValue(null, false);
77+
78+
// Act
79+
wrapper.Log("First message"); // This should cause a reset
80+
bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
81+
82+
SystemWrapper.ClearOutputResetFlag();
83+
bool afterClear = (bool)_outputResetPerformedField.GetValue(null);
84+
85+
wrapper.Log("After clear"); // This should cause another reset
86+
bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
87+
88+
// Assert
89+
Assert.True(afterFirstLog, "Flag should be set after first log");
90+
Assert.False(afterClear, "Flag should be cleared after ClearOutputResetFlag");
91+
Assert.True(afterSecondLog, "Flag should be set again after second log");
92+
}
93+
94+
[Fact]
95+
public void Log_InTestMode_WritesToTestOutput()
96+
{
97+
// Arrange
98+
var wrapper = new SystemWrapper(_mockEnvironment);
99+
SystemWrapper.SetOut(_testWriter);
100+
var message = "Test message";
101+
102+
// Act
103+
wrapper.Log(message);
104+
105+
// Assert
106+
Assert.Equal(message, _testWriter.ToString());
107+
}
108+
109+
[Fact]
110+
public void LogLine_InTestMode_WritesToTestOutput()
111+
{
112+
// Arrange
113+
var wrapper = new SystemWrapper(_mockEnvironment);
114+
SystemWrapper.SetOut(_testWriter);
115+
var message = "Test line";
116+
117+
// Act
118+
wrapper.LogLine(message);
119+
120+
// Assert
121+
Assert.Equal(message + Environment.NewLine, _testWriter.ToString());
122+
}
123+
124+
[Fact]
125+
public void ResetTestMode_ResetsTestState()
126+
{
127+
// Arrange
128+
var wrapper = new SystemWrapper(_mockEnvironment);
129+
SystemWrapper.SetOut(_testWriter);
130+
var message = "This should go to console";
131+
132+
// Act
133+
SystemWrapper.ResetTestMode();
134+
135+
// Can't directly test that this goes to console, but we can verify
136+
// it doesn't go to the test writer
137+
wrapper.Log(message);
138+
139+
// Assert
140+
Assert.Equal("", _testWriter.ToString());
141+
}
142+
143+
[Fact]
144+
public void SetOut_EnablesTestMode()
145+
{
146+
// Arrange
147+
var wrapper = new SystemWrapper(_mockEnvironment);
148+
var message = "Test output";
149+
150+
// Act
151+
SystemWrapper.SetOut(_testWriter);
152+
wrapper.Log(message);
153+
154+
// Assert
155+
Assert.Equal(message, _testWriter.ToString());
156+
}
157+
158+
[Fact]
159+
public void Log_InTestMode_DoesNotCallResetConsoleOutput()
160+
{
161+
// Arrange
162+
var wrapper = new SystemWrapper(_mockEnvironment);
163+
SystemWrapper.SetOut(_testWriter);
164+
var message1 = "First test message";
165+
var message2 = "Second test message";
166+
167+
// Act
168+
wrapper.Log(message1);
169+
wrapper.Log(message2);
170+
171+
// Assert
172+
Assert.Equal(message1 + message2, _testWriter.ToString());
173+
}
174+
175+
[Fact]
176+
public void Log_AfterClearingFlag_ResetsOutputAgain()
177+
{
178+
// Arrange
179+
var wrapper = new SystemWrapper(_mockEnvironment);
180+
_outputResetPerformedField.SetValue(null, false);
181+
182+
// Act
183+
wrapper.Log("First message"); // Should reset output
184+
bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
185+
186+
SystemWrapper.ClearOutputResetFlag();
187+
bool afterClear = (bool)_outputResetPerformedField.GetValue(null);
188+
189+
wrapper.Log("Second message"); // Should reset again
190+
bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
191+
192+
// Assert
193+
Assert.True(afterFirstLog, "Flag should be set after first log");
194+
Assert.False(afterClear, "Flag should be reset after clearing");
195+
Assert.True(afterSecondLog, "Flag should be set after second log");
196+
}
197+
198+
public void Dispose()
199+
{
200+
_testWriter?.Dispose();
201+
SystemWrapper.ResetTestMode();
202+
_outputResetPerformedField.SetValue(null, false);
203+
}
204+
}

0 commit comments

Comments
 (0)