redirect stdout/stderr to improve secret

lisa/tests/test_secret.py

@@ -59,6 +59,14 @@ class SecretTestCase(TestCase):
             log.info("t1t2 2")
         self.assertListEqual(["INFO:LISA:** 2"], cm.output)
 
+    def test_stdout(self) -> None:
+        add_secret("t1", sub="*")
+        add_secret("t1t2", sub="**")
+
+        with self.assertLogs("LISA") as cm:
+            print("t1t2 2")
+        self.assertListEqual(["INFO:LISA.stdout:** 2"], cm.output)
+
     def test_log_with_args(self) -> None:
         log = get_logger()
         add_secret("t1")

lisa/util/logger.py

@@ -1,7 +1,8 @@
 import logging
+import sys
 import time
 from functools import partial
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Optional, TextIO, Union, cast
 
 from lisa.secret import mask
 from lisa.util import LisaException
@@ -81,6 +82,26 @@ class Logger(logging.Logger):
             self.warn(message)
 
 
+class LogWriter(object):
+    def __init__(self, logger: Logger, level: int):
+        self._level = level
+        self._log = logger
+        self._buffer: str = ""
+
+    def write(self, message: str) -> None:
+        self._buffer = "".join([self._buffer, message])
+        if "\n" in message or "\r" in message:
+            self.flush()
+
+    def flush(self) -> None:
+        if len(self._buffer) > 0:
+            self._log.log(self._level, self._buffer.strip("\r\n"))
+            self._buffer = ""
+
+    def close(self) -> None:
+        self.flush()
+
+
 _get_root_logger = partial(logging.getLogger, DEFAULT_LOG_NAME)
 
 _format = "%(asctime)s.%(msecs)03d[%(levelname)-.1s]%(name)s %(message)s"
@@ -97,6 +118,11 @@ def init_loggger() -> None:
     logging.Formatter.converter = time.gmtime
     logging.setLoggerClass(Logger)
 
+    stdout_logger = get_logger("stdout")
+    stderr_logger = get_logger("stderr")
+    sys.stdout = cast(TextIO, LogWriter(stdout_logger, logging.INFO))
+    sys.stderr = cast(TextIO, LogWriter(stderr_logger, logging.ERROR))
+
 
 def set_log_file(path: str) -> None:
     root_logger = _get_root_logger()

lisa/util/process.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Dict, Optional, Type
 
 from spur.errors import NoSuchCommandError  # type: ignore
 
-from lisa.util.logger import Logger, get_logger
+from lisa.util.logger import Logger, LogWriter, get_logger
 from lisa.util.perf_timer import create_timer
 from lisa.util.shell import Shell
 
@@ -28,24 +28,6 @@ class ExecutableResult:
         return self.stdout
 
 
-class LogWriter:
-    def __init__(self, level: int, log: Logger) -> None:
-        self._level = level
-        self._log = log
-        self._buffer: str = ""
-
-    def write(self, message: str) -> None:
-        if message == "\n":
-            self._log.log(self._level, self._buffer)
-            self._buffer = ""
-        else:
-            self._buffer = "".join([self._buffer, message])
-
-    def close(self) -> None:
-        if len(self._buffer) > 0:
-            self._log.log(self._level, self._buffer)
-
-
 class Process:
     def __init__(
         self,
@@ -81,10 +63,10 @@ class Process:
             stderr_level = stdout_level
 
         self.stdout_writer = LogWriter(
-            stdout_level, get_logger("stdout", parent=self._log)
+            logger=get_logger("stdout", parent=self._log), level=stdout_level
         )
         self.stderr_writer = LogWriter(
-            stderr_level, get_logger("stderr", parent=self._log)
+            logger=get_logger("stderr", parent=self._log), level=stderr_level
         )
 
         # command may be Path object, convert it to str