diff options
| -rw-r--r-- | mitmproxy/addons/script.py | 16 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_script.py | 27 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/print.py | 2 |
3 files changed, 38 insertions, 7 deletions
diff --git a/mitmproxy/addons/script.py b/mitmproxy/addons/script.py index 4d893f1c..3b7a4f79 100644 --- a/mitmproxy/addons/script.py +++ b/mitmproxy/addons/script.py @@ -65,14 +65,28 @@ def cut_traceback(tb, func_name): return tb +class StreamLog: + """ + A class for redirecting output using contextlib. + """ + def __init__(self, log): + self.log = log + + def write(self, buf): + if buf.strip(): + self.log(buf) + + @contextlib.contextmanager def scriptenv(path, args): oldargs = sys.argv sys.argv = [path] + args script_dir = os.path.dirname(os.path.abspath(path)) sys.path.append(script_dir) + stdout_replacement = StreamLog(ctx.log.warn) try: - yield + with contextlib.redirect_stdout(stdout_replacement): + yield except SystemExit as v: ctx.log.error("Script exited with code %s" % v.code) except Exception: diff --git a/test/mitmproxy/addons/test_script.py b/test/mitmproxy/addons/test_script.py index 16827488..84d36e2a 100644 --- a/test/mitmproxy/addons/test_script.py +++ b/test/mitmproxy/addons/test_script.py @@ -5,6 +5,7 @@ import re import watchdog.events import pytest +from unittest import mock from mitmproxy.test import tflow from mitmproxy.test import tutils from mitmproxy.test import taddons @@ -97,12 +98,26 @@ class TestParseCommand: def test_load_script(): - ns = script.load_script( - tutils.test_data.path( - "mitmproxy/data/addonscripts/recorder.py" - ), [] - ) - assert ns.start + with taddons.context(): + ns = script.load_script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/recorder.py" + ), [] + ) + assert ns.start + + +def test_script_print_stdout(): + with taddons.context() as tctx: + with mock.patch('mitmproxy.ctx.log.warn') as mock_warn: + with script.scriptenv("path", []): + ns = script.load_script( + tutils.test_data.path( + "mitmproxy/data/addonscripts/print.py" + ), [] + ) + ns.start(tctx.options) + mock_warn.assert_called_once_with("stdoutprint") class TestScript: diff --git a/test/mitmproxy/data/addonscripts/print.py b/test/mitmproxy/data/addonscripts/print.py new file mode 100644 index 00000000..fdfcc3d9 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/print.py @@ -0,0 +1,2 @@ +def start(opts): + print("stdoutprint") |