addons.streambodies: 100% test coverage

2 files changed, 12 insertions, 14 deletions

diff --git a/mitmproxy/addons/streambodies.py b/mitmproxy/addons/streambodies.py
index 88cb74f3..3c2a153b 100644
--- a/mitmproxy/addons/streambodies.py
+++ b/mitmproxy/addons/streambodies.py
@@ -17,7 +17,7 @@ class StreamBodies:
                 expected_size = http1.expected_http_body_size(
                     f.request, f.response if not is_request else None
                 )
-            except exceptions.HTTPException:
+            except exceptions.HttpException:
                 f.reply.kill()
                 return
             if expected_size and not r.raw_content and not (0 <= expected_size <= self.max_size):
diff --git a/test/mitmproxy/addons/test_streambodies.py b/test/mitmproxy/addons/test_streambodies.py
index afd2cc51..b982c66d 100644
--- a/test/mitmproxy/addons/test_streambodies.py
+++ b/test/mitmproxy/addons/test_streambodies.py
@@ -1,30 +1,28 @@
 from mitmproxy.test import tflow
-
-from .. import mastertest
-from mitmproxy import master
-from mitmproxy import options
-from mitmproxy import proxy
+from mitmproxy.test import taddons
 
 from mitmproxy.addons import streambodies
 
 
-class TestStreamBodies(mastertest.MasterTest):
-    def test_simple(self):
-        o = options.Options(stream_large_bodies = 10)
-        m = master.Master(o, proxy.DummyServer())
-        sa = streambodies.StreamBodies()
-        m.addons.add(sa)
+def test_simple():
+    sa = streambodies.StreamBodies()
+    with taddons.context() as tctx:
+        tctx.configure(sa, stream_large_bodies = 10)
 
         f = tflow.tflow()
         f.request.content = b""
         f.request.headers["Content-Length"] = "1024"
         assert not f.request.stream
-        m.requestheaders(f)
+        sa.requestheaders(f)
         assert f.request.stream
 
         f = tflow.tflow(resp=True)
         f.response.content = b""
         f.response.headers["Content-Length"] = "1024"
         assert not f.response.stream
-        m.responseheaders(f)
+        sa.responseheaders(f)
         assert f.response.stream
+
+        f = tflow.tflow(resp=True)
+        f.response.headers["content-length"] = "invalid"
+        tctx.cycle(sa, f)