simplify termlog outfile handling

2 files changed, 28 insertions, 22 deletions

diff --git a/mitmproxy/addons/termlog.py b/mitmproxy/addons/termlog.py
index 8ee366cd..5fdb6245 100644
--- a/mitmproxy/addons/termlog.py
+++ b/mitmproxy/addons/termlog.py
@@ -7,24 +7,21 @@ from mitmproxy import log
 class TermLog:
     def __init__(self, outfile=None):
         self.options = None
-        self.outfile = outfile or sys.stdout
+        self.outfile = outfile
 
     def configure(self, options, updated):
         self.options = options
 
     def log(self, e):
-        def determine_outfile():
-            if log.log_tier(e.level) == log.log_tier("error"):
-                self.outfile = sys.stderr
-            else:
-                self.outfile = sys.stdout
-
-        determine_outfile()
+        if log.log_tier(e.level) == log.log_tier("error"):
+            outfile = self.outfile or sys.stderr
+        else:
+            outfile = self.outfile or sys.stdout
 
         if self.options.verbosity >= log.log_tier(e.level):
             click.secho(
                 e.msg,
-                file=self.outfile,
+                file=outfile,
                 fg=dict(error="red", warn="yellow").get(e.level),
                 dim=(e.level == "debug"),
                 err=(e.level == "error")
diff --git a/test/mitmproxy/addons/test_termlog.py b/test/mitmproxy/addons/test_termlog.py
index 9d49a8c0..e21aecec 100644
--- a/test/mitmproxy/addons/test_termlog.py
+++ b/test/mitmproxy/addons/test_termlog.py
@@ -1,18 +1,27 @@
+import sys
+import pytest
+
 from mitmproxy.addons import termlog
 from mitmproxy import log
-from mitmproxy.tools import dump
+from mitmproxy.tools.dump import Options
+from mitmproxy.test import taddons
 
 
 class TestTermLog:
-    def test_simple(self, capsys):
-        t = termlog.TermLog()
-        t.configure(dump.Options(verbosity = 2), set([]))
-        t.log(log.LogEntry("one", "info"))
-        t.log(log.LogEntry("two", "debug"))
-        t.log(log.LogEntry("three", "warn"))
-        t.log(log.LogEntry("four", "error"))
-        out, err = capsys.readouterr()
-        assert "one" in out
-        assert "two" not in out
-        assert "three" in out
-        assert "four" in err
+    @pytest.mark.usefixtures('capfd')
+    @pytest.mark.parametrize('outfile, expected_out, expected_err', [
+        (None, 'one\nthree\n', 'four\n'),
+        (sys.stdout, 'one\nthree\nfour\n', ''),
+        (sys.stderr, '', 'one\nthree\nfour\n'),
+    ])
+    def test_output(self, outfile, expected_out, expected_err, capfd):
+        t = termlog.TermLog(outfile=outfile)
+        with taddons.context(options=Options(verbosity=2)) as tctx:
+            tctx.configure(t)
+            t.log(log.LogEntry("one", "info"))
+            t.log(log.LogEntry("two", "debug"))
+            t.log(log.LogEntry("three", "warn"))
+            t.log(log.LogEntry("four", "error"))
+        out, err = capfd.readouterr()
+        assert out == expected_out
+        assert err == expected_err