Merge pull request #1708 from cortesi/taddons3

Test suite improvements

2 files changed, 62 insertions, 14 deletions

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)
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index de752c72..5243e9d4 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -9,27 +9,68 @@ from mitmproxy.test import taddons
 
 class Options(options.Options):
     def __init__(
-        self, *,
+        self,
+        *,
         filter=None,
         order=None,
         order_reversed=False,
+        focus_follow=False,
         **kwargs
     ):
         self.filter = filter
         self.order = order
         self.order_reversed = order_reversed
+        self.focus_follow = focus_follow
         super().__init__(**kwargs)
 
 
+def test_order_refresh():
+    v = view.View()
+    sargs = []
+
+    def save(*args, **kwargs):
+        sargs.extend([args, kwargs])
+
+    v.sig_refresh.connect(save)
+
+    tf = tflow.tflow(resp=True)
+    with taddons.context(options=Options()) as tctx:
+        tctx.configure(v, order="time")
+        v.add(tf)
+        tf.request.timestamp_start = 1
+        assert not sargs
+        v.update(tf)
+        assert sargs
+
+
+def test_order_generators():
+    v = view.View()
+    tf = tflow.tflow(resp=True)
+
+    rs = view.OrderRequestStart(v)
+    assert rs.generate(tf) == 0
+
+    rm = view.OrderRequestMethod(v)
+    assert rm.generate(tf) == tf.request.method
+
+    ru = view.OrderRequestURL(v)
+    assert ru.generate(tf) == tf.request.url
+
+    sz = view.OrderKeySize(v)
+    assert sz.generate(tf) == len(tf.request.raw_content) + len(tf.response.raw_content)
+
+
 def test_simple():
     v = view.View()
     f = tflow.tflow()
+    assert v.store_count() == 0
     f.request.timestamp_start = 1
     v.request(f)
     assert list(v) == [f]
     v.request(f)
     assert list(v) == [f]
     assert len(v._store) == 1
+    assert v.store_count() == 1
 
     f2 = tflow.tflow()
     f2.request.timestamp_start = 3
@@ -39,6 +80,10 @@ def test_simple():
     assert list(v) == [f, f2]
     assert len(v._store) == 2
 
+    assert v.inbounds(0)
+    assert not v.inbounds(-1)
+    assert not v.inbounds(100)
+
     f3 = tflow.tflow()
     f3.request.timestamp_start = 2
     v.request(f3)
@@ -289,3 +334,8 @@ def test_configure():
         tutils.raises("unknown flow order", tctx.configure, v, order="no")
 
         tctx.configure(v, order_reversed=True)
+
+        tctx.configure(v, order=None)
+
+        tctx.configure(v, focus_follow=True)
+        assert v.focus_follow