From b5f31daab0109f01fe865a7cbd67ac7a9857886a Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Tue, 30 Mar 2021 13:44:46 -0700
Subject: [PATCH] ui: Use sys.stdout and sys.stderr dynamically (#5729)

Before, we used to set it during initialization. But this
made it hard to test using capsys as capsys changes the
reference to sys.stdout and sys.stderr with the mocked one
during the test. But since we are using dvc.ui globally and
setting it to a real sys.stdout and sys.stderr before capsys
comes into effect, testing was not possible with it.

Now, by using property hack, we will be using the values of
sys.stdout and sys.stderr during the time of write, so the mocking
would work and would be easier to test.
---
 dvc/ui/__init__.py            | 22 ++++++++++++++--------
 tests/conftest.py             |  7 +++++++
 tests/unit/ui/__init__.py     |  0
 tests/unit/ui/test_console.py | 28 ++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 8 deletions(-)
 create mode 100644 tests/unit/ui/__init__.py
 create mode 100644 tests/unit/ui/test_console.py

diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py
index 1b6cd6c8b..164a859ad 100644
--- a/dvc/ui/__init__.py
+++ b/dvc/ui/__init__.py
@@ -29,13 +29,19 @@ class Console:
         error: TextIO = None,
         enable: bool = False,
     ) -> None:
-        self._input: TextIO = sys.stdin
-        self._output: TextIO = output or sys.stdout
-        self._error: TextIO = error or sys.stderr
-
+        self._output: Optional[TextIO] = output
+        self._error: Optional[TextIO] = error
         self.formatter: Formatter = formatter or Formatter()
         self._enabled: bool = enable
 
+    @property
+    def output(self) -> TextIO:
+        return self._output or sys.stdout
+
+    @property
+    def error_output(self) -> TextIO:
+        return self._error or sys.stderr
+
     def enable(self):
         self._enabled = True
 
@@ -61,7 +67,7 @@ class Console:
             style=style,
             sep=sep,
             end=end,
-            file=self._error,
+            file=self.error_output,
             flush=flush,
         )
 
@@ -77,12 +83,12 @@ class Console:
         if not self._enabled:
             return
 
-        file = file or self._output
+        file = file or self.output
         values = (self.formatter.format(obj, style=style) for obj in objects)
         return print(*values, sep=sep, end=end, file=file, flush=flush)
 
     def progress(self, *args, **kwargs) -> Tqdm:
-        kwargs.setdefault("file", self._error)
+        kwargs.setdefault("file", self.error_output)
         return Tqdm(*args, **kwargs)
 
     def prompt(
@@ -125,7 +131,7 @@ class Console:
         """rich_console is only set to stdout for now."""
         from rich import console
 
-        return console.Console(file=self._output)
+        return console.Console(file=self.output)
 
     def rich_table(self, pager: bool = True):
         pass
diff --git a/tests/conftest.py b/tests/conftest.py
index e5d906fa1..4a000c5b4 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -43,6 +43,13 @@ def reset_loglevel(request, caplog):
         yield
 
 
+@pytest.fixture(autouse=True)
+def enable_ui():
+    from dvc.ui import ui
+
+    ui.enable()
+
+
 @pytest.fixture(scope="session", autouse=True)
 def _close_pools():
     from dvc.fs.pool import close_pools
diff --git a/tests/unit/ui/__init__.py b/tests/unit/ui/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/unit/ui/test_console.py b/tests/unit/ui/test_console.py
new file mode 100644
index 000000000..9e3680493
--- /dev/null
+++ b/tests/unit/ui/test_console.py
@@ -0,0 +1,28 @@
+from pytest import CaptureFixture
+
+from dvc.ui import Console
+
+
+def test_write(capsys: CaptureFixture[str]):
+    """Test that ui.write works."""
+    console = Console(enable=True)
+    message = "hello world"
+    console.write(message)
+    console.error_write(message)
+
+    captured = capsys.readouterr()
+    assert captured.out == f"{message}\n"
+    assert captured.err == f"{message}\n"
+
+
+def test_capsys_works(capsys: CaptureFixture[str]):
+    """Sanity check that capsys can capture outputs from a global ui."""
+    from dvc.ui import ui
+
+    message = "hello world"
+    ui.write(message)
+    ui.error_write(message)
+
+    captured = capsys.readouterr()
+    assert captured.out == f"{message}\n"
+    assert captured.err == f"{message}\n"