addons.view.focus: Better handling of view refresh

When we refresh and our current focus goes out of scope, we set the focus to
the element nearest the old focus.

2 files changed, 6 insertions, 4 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index b5682b7a..d8d6e853 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -199,18 +199,20 @@ class Focus:
         if self.focusflow:
             return self.view.index(self.focusflow)
 
+    def _nearest(self, f, v):
+        return min(v.bisect(f), len(v)-1)
+
     def _sig_remove(self, view, flow):
         if len(view) == 0:
             self.focusflow = None
         elif flow is self.focusflow:
-            idx = min(view.bisect(self.focusflow), len(view)-1)
-            self.focusflow = view[idx]
+            self.focusflow = view[self._nearest(self.focusflow, view)]
 
     def _sig_refresh(self, view):
         if len(view) == 0:
             self.focusflow = None
         elif self.focusflow not in view:
-            self.focusflow = view[0]
+            self.focusflow = view[self._nearest(self.focusflow, view)]
 
     def _sig_add(self, view, flow):
         # We only have to act if we don't have a focus element
diff --git a/test/mitmproxy/addons/test_view.py b/test/mitmproxy/addons/test_view.py
index 1ab7e8f4..56372749 100644
--- a/test/mitmproxy/addons/test_view.py
+++ b/test/mitmproxy/addons/test_view.py
@@ -228,7 +228,7 @@ def test_focus():
 
     filt = flowfilter.parse("~m get")
     v.set_filter(filt)
-    assert f.index == 0
+    assert f.index == 2
 
     filt = flowfilter.parse("~m oink")
     v.set_filter(filt)