From e6419b939df7c6fe2b026cb2ad8ead7b096cb34a Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Wed, 15 Mar 2023 18:23:26 +0545
Subject: [PATCH] do not disable colorama in tests (#9187)

skip disabling colorama in tests

It might affect pytest as it wraps sys.stderr/sys.stdout for capsys fixture.
This is a temporary hack, as we might get rid of colorama soon.
---
 dvc/ui/__init__.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/dvc/ui/__init__.py b/dvc/ui/__init__.py
index 5e81a2e75c..0fc60bceb1 100644
--- a/dvc/ui/__init__.py
+++ b/dvc/ui/__init__.py
@@ -1,4 +1,4 @@
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -132,9 +132,12 @@ def write_json(  # noqa: PLR0913
             sort_keys=sort_keys,
         )
         if not highlight:
+            import os
+
             # we don't need colorama to try to strip ansi codes
             # when highlighting is disabled
-            with disable_colorama():
+            ctx = nullcontext() if "DVC_TEST" in os.environ else disable_colorama()
+            with ctx:
                 return self.write(json.text, stderr=stderr)
         return self.rich_print(json, stderr=stderr, soft_wrap=True)