Redesign keepserving

- Instead of listening for a pseudo-event, we periodically check whether client
replay, server replay or file reading is active.
- Adjust server replay not to
use tick.
- Adjust readfile to expose a command to check whether reading is in progress.

4 files changed, 43 insertions, 22 deletions

diff --git a/test/mitmproxy/addons/test_keepserving.py b/test/mitmproxy/addons/test_keepserving.py
index 5eafa792..01b0d09c 100644
--- a/test/mitmproxy/addons/test_keepserving.py
+++ b/test/mitmproxy/addons/test_keepserving.py
@@ -3,12 +3,48 @@ import pytest
 
 from mitmproxy.addons import keepserving
 from mitmproxy.test import taddons
+from mitmproxy import command
+
+
+class Dummy:
+    def __init__(self, val: bool):
+        self.val = val
+
+    def load(self, loader):
+        loader.add_option("client_replay", bool, self.val, "test")
+        loader.add_option("server_replay", bool, self.val, "test")
+        loader.add_option("rfile", bool, self.val, "test")
+
+    @command.command("readfile.reading")
+    def readfile(self) -> bool:
+        return self.val
+
+    @command.command("replay.client.count")
+    def creplay(self) -> int:
+        return 1 if self.val else 0
+
+    @command.command("replay.server.count")
+    def sreplay(self) -> int:
+        return 1 if self.val else 0
+
+
+class TKS(keepserving.KeepServing):
+    _is_shutdown = False
+
+    def shutdown(self):
+        self.is_shutdown = True
 
 
 @pytest.mark.asyncio
 async def test_keepserving():
-    ks = keepserving.KeepServing()
+    ks = TKS()
+    d = Dummy(True)
     with taddons.context(ks) as tctx:
-        ks.event_processing_complete()
-        asyncio.sleep(0.1)
-        assert tctx.master.should_exit.is_set()
+        tctx.master.addons.add(d)
+        ks.running()
+        assert ks.keepgoing()
+
+        d.val = False
+        assert not ks.keepgoing()
+        await asyncio.sleep(0.3)
+        assert ks.is_shutdown
diff --git a/test/mitmproxy/addons/test_readfile.py b/test/mitmproxy/addons/test_readfile.py
index d22382a8..62f4d917 100644
--- a/test/mitmproxy/addons/test_readfile.py
+++ b/test/mitmproxy/addons/test_readfile.py
@@ -51,6 +51,8 @@ class TestReadFile:
     async def test_read(self, tmpdir, data, corrupt_data):
         rf = readfile.ReadFile()
         with taddons.context(rf) as tctx:
+            assert not rf.reading()
+
             tf = tmpdir.join("tfile")
 
             with asynctest.patch('mitmproxy.master.Master.load_flow') as mck:
diff --git a/test/mitmproxy/addons/test_serverplayback.py b/test/mitmproxy/addons/test_serverplayback.py
index 0bc28ac8..c6a0c1f4 100644
--- a/test/mitmproxy/addons/test_serverplayback.py
+++ b/test/mitmproxy/addons/test_serverplayback.py
@@ -39,16 +39,6 @@ def test_config(tmpdir):
             tctx.configure(s, server_replay=[str(tmpdir)])
 
 
-def test_tick():
-    s = serverplayback.ServerPlayback()
-    with taddons.context(s) as tctx:
-        s.stop = True
-        s.final_flow = tflow.tflow()
-        s.final_flow.live = False
-        s.tick()
-        assert tctx.master.has_event("processing_complete")
-
-
 def test_server_playback():
     sp = serverplayback.ServerPlayback()
     with taddons.context(sp) as tctx:
@@ -349,14 +339,6 @@ def test_server_playback_full():
         s.request(tf)
         assert not tf.response
 
-        assert not s.stop
-        s.tick()
-        assert not s.stop
-
-        tf = tflow.tflow()
-        s.request(tflow.tflow())
-        assert s.stop
-
 
 def test_server_playback_kill():
     s = serverplayback.ServerPlayback()
diff --git a/test/mitmproxy/tools/console/test_statusbar.py b/test/mitmproxy/tools/console/test_statusbar.py
index 108f238e..f1cc67d4 100644
--- a/test/mitmproxy/tools/console/test_statusbar.py
+++ b/test/mitmproxy/tools/console/test_statusbar.py
@@ -30,6 +30,7 @@ def test_statusbar(monkeypatch):
     m.options.update(view_order='url', console_focus_follow=True)
     monkeypatch.setattr(m.addons.get("clientplayback"), "count", lambda: 42)
     monkeypatch.setattr(m.addons.get("serverplayback"), "count", lambda: 42)
+    monkeypatch.setattr(statusbar.StatusBar, "refresh", lambda x: None)
 
     bar = statusbar.StatusBar(m)  # this already causes a redraw
     assert bar.ib._w