addons.View: hook up signals

2 files changed, 81 insertions, 4 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index f79dd64d..76f82ed6 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -36,13 +36,11 @@ class View(collections.Sequence):
         self.filter = matchall
         self.order_key = key_request_start
         self.order_reverse = False
-        self._view = sortedcontainers.SortedListWithKey(
-            key = self.order_key
-        )
+        self._view = sortedcontainers.SortedListWithKey(key = self.order_key)
 
         # These signals broadcast events that affect the view. That is, an
         # update to a flow in the store but not in the view does not trigger a
-        # signal.
+        # signal. All signals are called after the view has been updated.
         self.sig_update = blinker.Signal()
         self.sig_add = blinker.Signal()
         self.sig_remove = blinker.Signal()
@@ -75,6 +73,7 @@ class View(collections.Sequence):
         for i in self._store.values():
             if self.filter(i):
                 self._view.add(i)
+        self.sig_refresh.send(self)
 
     def clear(self):
         """
@@ -82,6 +81,7 @@ class View(collections.Sequence):
         """
         self._state.clear()
         self._view.clear()
+        self.sig_refresh.send(self)
 
     def add(self, f: flow.Flow):
         """
@@ -92,6 +92,7 @@ class View(collections.Sequence):
             self._store[f.id] = f
             if self.filter(f):
                 self._view.add(f)
+                self.sig_add.send(self, flow=f)
 
     def update(self, f: flow.Flow):
         """
@@ -101,10 +102,15 @@ class View(collections.Sequence):
             if self.filter(f):
                 if f not in self._view:
                     self._view.add(f)
+                    self.sig_add.send(self, flow=f)
+                else:
+                    self.sig_update.send(self, flow=f)
             else:
                 try:
                     self._view.remove(f)
+                    self.sig_remove.send(self, flow=f)
                 except ValueError:
+                    # The value was not in the view
                     pass
 
     # Event handlers
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 3e01d71f..cff99621 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -1,5 +1,6 @@
 from mitmproxy.addons import view
 from mitmproxy import flowfilter
+
 from .. import tutils
 
 
@@ -90,3 +91,73 @@ def test_update():
 
     v.update(f)
     assert f in v
+
+
+class Record:
+    def __init__(self):
+        self.calls = []
+
+    def __bool__(self):
+        return bool(self.calls)
+
+    def __repr__(self):
+        return repr(self.calls)
+
+    def __call__(self, *args, **kwargs):
+        self.calls.append((args, kwargs))
+
+
+def test_signals():
+    v = view.View()
+    rec_add = Record()
+    rec_update = Record()
+    rec_remove = Record()
+    rec_refresh = Record()
+
+    def clearrec():
+        rec_add.calls = []
+        rec_update.calls = []
+        rec_remove.calls = []
+        rec_refresh.calls = []
+
+    v.sig_add.connect(rec_add)
+    v.sig_update.connect(rec_update)
+    v.sig_remove.connect(rec_remove)
+    v.sig_refresh.connect(rec_refresh)
+
+    assert not any([rec_add, rec_update, rec_remove, rec_refresh])
+
+    # Simple add
+    v.add(tft())
+    assert rec_add
+    assert not any([rec_update, rec_remove, rec_refresh])
+
+    # Filter change triggers refresh
+    clearrec()
+    v.set_filter(flowfilter.parse("~m put"))
+    assert rec_refresh
+    assert not any([rec_update, rec_add, rec_remove])
+
+    v.set_filter(flowfilter.parse("~m get"))
+
+    # An update that results in a flow being added to the view
+    clearrec()
+    v[0].request.method = "PUT"
+    v.update(v[0])
+    assert rec_remove
+    assert not any([rec_update, rec_refresh, rec_add])
+
+    # An update that does not affect the view just sends update
+    v.set_filter(flowfilter.parse("~m put"))
+    clearrec()
+    v.update(v[0])
+    assert rec_update
+    assert not any([rec_remove, rec_refresh, rec_add])
+
+    # An update for a flow in state but not view does not do anything
+    f = v[0]
+    v.set_filter(flowfilter.parse("~m get"))
+    assert not len(v)
+    clearrec()
+    v.update(f)
+    assert not any([rec_add, rec_update, rec_remove, rec_refresh])