@@ -22,41 +22,77 @@ namespace AWS.Lambda.Powertools.Common;
22
22
public class ConsoleWrapper : IConsoleWrapper
23
23
{
24
24
private static bool _override ;
25
+ private static TextWriter _testOutputStream ;
26
+ private static bool _outputResetPerformed = false ;
27
+ private static bool _inTestMode = false ;
25
28
26
29
/// <inheritdoc />
27
30
public void WriteLine ( string message )
28
31
{
29
- OverrideLambdaLogger ( ) ;
30
- Console . WriteLine ( message ) ;
32
+ if ( _inTestMode && _testOutputStream != null )
33
+ {
34
+ _testOutputStream . WriteLine ( message ) ;
35
+ }
36
+ else
37
+ {
38
+ EnsureConsoleOutputOnce ( ) ;
39
+ Console . WriteLine ( message ) ;
40
+ }
31
41
}
32
42
33
43
/// <inheritdoc />
34
44
public void Debug ( string message )
35
45
{
36
- OverrideLambdaLogger ( ) ;
37
- System . Diagnostics . Debug . WriteLine ( message ) ;
46
+ if ( _inTestMode && _testOutputStream != null )
47
+ {
48
+ _testOutputStream . WriteLine ( message ) ;
49
+ }
50
+ else
51
+ {
52
+ EnsureConsoleOutputOnce ( ) ;
53
+ System . Diagnostics . Debug . WriteLine ( message ) ;
54
+ }
38
55
}
39
56
40
57
/// <inheritdoc />
41
58
public void Error ( string message )
42
59
{
43
- if ( ! _override )
60
+ if ( _inTestMode && _testOutputStream != null )
44
61
{
45
- var errordOutput = new StreamWriter ( Console . OpenStandardError ( ) ) ;
46
- errordOutput . AutoFlush = true ;
47
- Console . SetError ( errordOutput ) ;
62
+ _testOutputStream . WriteLine ( message ) ;
63
+ }
64
+ else
65
+ {
66
+ if ( ! _override )
67
+ {
68
+ var errordOutput = new StreamWriter ( Console . OpenStandardError ( ) ) ;
69
+ errordOutput . AutoFlush = true ;
70
+ Console . SetError ( errordOutput ) ;
71
+ }
72
+ Console . Error . WriteLine ( message ) ;
48
73
}
49
-
50
- Console . Error . WriteLine ( message ) ;
51
74
}
52
75
53
- internal static void SetOut ( StringWriter consoleOut )
76
+ /// <summary>
77
+ /// Set the ConsoleWrapper to use a different TextWriter
78
+ /// This is useful for unit tests where you want to capture the output
79
+ /// </summary>
80
+ public static void SetOut ( TextWriter consoleOut )
54
81
{
82
+ _testOutputStream = consoleOut ;
83
+ _inTestMode = true ;
55
84
_override = true ;
56
85
Console . SetOut ( consoleOut ) ;
57
86
}
58
87
59
- private void OverrideLambdaLogger ( )
88
+ private static void EnsureConsoleOutputOnce ( )
89
+ {
90
+ if ( _outputResetPerformed ) return ;
91
+ OverrideLambdaLogger ( ) ;
92
+ _outputResetPerformed = true ;
93
+ }
94
+
95
+ private static void OverrideLambdaLogger ( )
60
96
{
61
97
if ( _override )
62
98
{
@@ -73,8 +109,22 @@ internal static void WriteLine(string logLevel, string message)
73
109
Console . WriteLine ( $ "{ DateTime . UtcNow : yyyy-MM-ddTHH:mm:ss.fffZ} \t { logLevel } \t { message } ") ;
74
110
}
75
111
112
+ /// <summary>
113
+ /// Reset the ConsoleWrapper to its original state
114
+ /// </summary>
76
115
public static void ResetForTest ( )
77
116
{
78
117
_override = false ;
118
+ _inTestMode = false ;
119
+ _testOutputStream = null ;
120
+ _outputResetPerformed = false ;
121
+ }
122
+
123
+ /// <summary>
124
+ /// Clear the output reset flag
125
+ /// </summary>
126
+ public static void ClearOutputResetFlag ( )
127
+ {
128
+ _outputResetPerformed = false ;
79
129
}
80
130
}
0 commit comments