@@ -27,6 +27,9 @@ namespace AWS.Lambda.Powertools.Common;
2727public 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}
0 commit comments