chore: Fix Console output when running on Top Level Statements

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

@@ -59,15 +59,4 @@ public interface ISystemWrapper
     /// </summary>
     /// <param name="type"></param>
     void SetExecutionEnvironment<T>(T type);
-
-    /// <summary>
-    /// Sets console output
-    /// Useful for testing and checking the console output
-    /// <code>
-    /// var consoleOut = new StringWriter();
-    /// SystemWrapper.Instance.SetOut(consoleOut);
-    /// </code>
-    /// </summary>
-    /// <param name="writeTo">The TextWriter instance where to write to</param>
-    void SetOut(TextWriter writeTo);
 }

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

@@ -27,6 +27,9 @@
 public class SystemWrapper : ISystemWrapper
 {
     private static IPowertoolsEnvironment _powertoolsEnvironment;
+    private static bool _inTestMode = false;
+    private static TextWriter _testOutputStream;
+    private static bool _outputResetPerformed = false;
 
     /// <summary>
     ///     The instance
@@ -41,13 +44,11 @@
         _powertoolsEnvironment = powertoolsEnvironment;
         _instance ??= this;
 
-        // Clear AWS SDK Console injected parameters StdOut and StdErr
-        var standardOutput = new StreamWriter(Console.OpenStandardOutput());
-        standardOutput.AutoFlush = true;
-        Console.SetOut(standardOutput);
-        var errordOutput = new StreamWriter(Console.OpenStandardError());
-        errordOutput.AutoFlush = true;
-        Console.SetError(errordOutput);
+        if (!_inTestMode)
+        {
+            // Clear AWS SDK Console injected parameters in production only
+            ResetConsoleOutput();
+        }
     }
 
     /// <summary>
@@ -72,7 +73,15 @@
     /// <param name="value">The value.</param>
     public void Log(string value)
     {
-        Console.Write(value);
+        if (_inTestMode && _testOutputStream != null)
+        {
+            _testOutputStream.Write(value);
+        }
+        else
+        {
+            EnsureConsoleOutputOnce();
+            Console.Write(value);
+        }
     }
 
     /// <summary>
@@ -81,7 +90,15 @@
     /// <param name="value">The value.</param>
     public void LogLine(string value)
     {
-        Console.WriteLine(value);
+        if (_inTestMode && _testOutputStream != null)
+        {
+            _testOutputStream.WriteLine(value);
+        }
+        else
+        {
+            EnsureConsoleOutputOnce();
+            Console.WriteLine(value);
+        }
     }
 
     /// <summary>
@@ -126,9 +143,20 @@
         SetEnvironmentVariable(envName, envValue.ToString());
     }
 
-    /// <inheritdoc />
-    public void SetOut(TextWriter writeTo)
+    /// <summary>
+    /// Sets console output
+    /// Useful for testing and checking the console output
+    /// <code>
+    /// var consoleOut = new StringWriter();
+    /// SystemWrapper.Instance.SetOut(consoleOut);
+    /// </code>
+    /// </summary>
+    /// <param name="writeTo">The TextWriter instance where to write to</param>
+
+    public static void SetOut(TextWriter writeTo)
     {
+        _testOutputStream = writeTo;
+        _inTestMode = true;
         Console.SetOut(writeTo);
     }
 
@@ -152,4 +180,33 @@
 
         return $"{Constants.FeatureContextIdentifier}/{assemblyName}";
     }
+
+    private static void EnsureConsoleOutputOnce()
+    {
+        if (_outputResetPerformed) return;
+        ResetConsoleOutput();
+        _outputResetPerformed = true;
+    }
+
+    private static void ResetConsoleOutput()
+    {
+        var standardOutput = new StreamWriter(Console.OpenStandardOutput());
+        standardOutput.AutoFlush = true;
+        Console.SetOut(standardOutput);
+        var errorOutput = new StreamWriter(Console.OpenStandardError());
+        errorOutput.AutoFlush = true;
+        Console.SetError(errorOutput);
+    }
+
+    public static void ClearOutputResetFlag()
+    {
+        _outputResetPerformed = false;
+    }
+
+    // For test cleanup
+    internal static void ResetTestMode()
+    {
+        _inTestMode = false;
+        _testOutputStream = null;
+    }
 }

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

@@ -0,0 +1,204 @@
+using System;
+using System.IO;
+using System.Reflection;
+using NSubstitute;
+using Xunit;
+
+namespace AWS.Lambda.Powertools.Common.Tests;
+
+[Collection("Sequential")]
+public class SystemWrapperTests : IDisposable
+{
+    private readonly IPowertoolsEnvironment _mockEnvironment;
+    private readonly StringWriter _testWriter;
+    private readonly FieldInfo _outputResetPerformedField;
+
+
+    public SystemWrapperTests()
+    {
+        _mockEnvironment = Substitute.For<IPowertoolsEnvironment>();
+        _testWriter = new StringWriter();
+
+        // Get access to private field for testing
+        _outputResetPerformedField = typeof(SystemWrapper).GetField("_outputResetPerformed", 
+            BindingFlags.NonPublic | BindingFlags.Static);
+
+        // Reset static state between tests
+        SystemWrapper.ResetTestMode();
+        _outputResetPerformedField.SetValue(null, false);
+    }
+
+    [Fact]
+    public void Log_InProductionMode_ResetsOutputOnce()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        var message1 = "First message";
+        var message2 = "Second message";
+        _outputResetPerformedField.SetValue(null, false);
+
+        // Act
+        wrapper.Log(message1);
+        bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
+        wrapper.Log(message2);
+        bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        // Assert
+        Assert.True(afterFirstLog, "Flag should be set after first log");
+        Assert.True(afterSecondLog, "Flag should remain set after second log");
+    }
+
+    [Fact]
+    public void LogLine_InProductionMode_ResetsOutputOnce()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        var message1 = "First line";
+        var message2 = "Second line";
+        _outputResetPerformedField.SetValue(null, false);
+
+        // Act
+        wrapper.LogLine(message1);
+        bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
+        wrapper.LogLine(message2);
+        bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        // Assert
+        Assert.True(afterFirstLog, "Flag should be set after first LogLine");
+        Assert.True(afterSecondLog, "Flag should remain set after second LogLine");
+    }
+
+    [Fact]
+    public void ClearOutputResetFlag_ResetsFlag_AllowsSubsequentReset()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        _outputResetPerformedField.SetValue(null, false);
+
+        // Act
+        wrapper.Log("First message"); // This should cause a reset
+        bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        SystemWrapper.ClearOutputResetFlag();
+        bool afterClear = (bool)_outputResetPerformedField.GetValue(null);
+
+        wrapper.Log("After clear"); // This should cause another reset
+        bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        // Assert
+        Assert.True(afterFirstLog, "Flag should be set after first log");
+        Assert.False(afterClear, "Flag should be cleared after ClearOutputResetFlag");
+        Assert.True(afterSecondLog, "Flag should be set again after second log");
+    }
+
+    [Fact]
+    public void Log_InTestMode_WritesToTestOutput()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        SystemWrapper.SetOut(_testWriter);
+        var message = "Test message";
+
+        // Act
+        wrapper.Log(message);
+
+        // Assert
+        Assert.Equal(message, _testWriter.ToString());
+    }
+
+    [Fact]
+    public void LogLine_InTestMode_WritesToTestOutput()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        SystemWrapper.SetOut(_testWriter);
+        var message = "Test line";
+
+        // Act
+        wrapper.LogLine(message);
+
+        // Assert
+        Assert.Equal(message + Environment.NewLine, _testWriter.ToString());
+    }
+
+    [Fact]
+    public void ResetTestMode_ResetsTestState()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        SystemWrapper.SetOut(_testWriter);
+        var message = "This should go to console";
+
+        // Act
+        SystemWrapper.ResetTestMode();
+
+        // Can't directly test that this goes to console, but we can verify
+        // it doesn't go to the test writer
+        wrapper.Log(message);
+
+        // Assert
+        Assert.Equal("", _testWriter.ToString());
+    }
+
+    [Fact]
+    public void SetOut_EnablesTestMode()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        var message = "Test output";
+
+        // Act
+        SystemWrapper.SetOut(_testWriter);
+        wrapper.Log(message);
+
+        // Assert
+        Assert.Equal(message, _testWriter.ToString());
+    }
+
+    [Fact]
+    public void Log_InTestMode_DoesNotCallResetConsoleOutput()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        SystemWrapper.SetOut(_testWriter);
+        var message1 = "First test message";
+        var message2 = "Second test message";
+
+        // Act
+        wrapper.Log(message1);
+        wrapper.Log(message2);
+
+        // Assert
+        Assert.Equal(message1 + message2, _testWriter.ToString());
+    }
+
+    [Fact]
+    public void Log_AfterClearingFlag_ResetsOutputAgain()
+    {
+        // Arrange
+        var wrapper = new SystemWrapper(_mockEnvironment);
+        _outputResetPerformedField.SetValue(null, false);
+
+        // Act
+        wrapper.Log("First message"); // Should reset output
+        bool afterFirstLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        SystemWrapper.ClearOutputResetFlag();
+        bool afterClear = (bool)_outputResetPerformedField.GetValue(null);
+
+        wrapper.Log("Second message"); // Should reset again
+        bool afterSecondLog = (bool)_outputResetPerformedField.GetValue(null);
+
+        // Assert
+        Assert.True(afterFirstLog, "Flag should be set after first log");
+        Assert.False(afterClear, "Flag should be reset after clearing");
+        Assert.True(afterSecondLog, "Flag should be set after second log");
+    }
+
+    public void Dispose()
+    {
+        _testWriter?.Dispose();
+        SystemWrapper.ResetTestMode();
+        _outputResetPerformedField.SetValue(null, false);
+    }
+}