diff options
| -rw-r--r-- | mitmproxy/addons/clientplayback.py | 28 | ||||
| -rw-r--r-- | mitmproxy/proxy/protocol/http_replay.py | 1 | ||||
| -rw-r--r-- | test/mitmproxy/addons/test_clientplayback.py | 7 |
3 files changed, 22 insertions, 14 deletions
diff --git a/mitmproxy/addons/clientplayback.py b/mitmproxy/addons/clientplayback.py index 0db6d336..9e012b67 100644 --- a/mitmproxy/addons/clientplayback.py +++ b/mitmproxy/addons/clientplayback.py @@ -9,15 +9,16 @@ import typing class ClientPlayback: def __init__(self): - self.flows = None + self.flows = [] # type: typing.List[flow.Flow] self.current_thread = None - self.has_replayed = False self.configured = False def count(self) -> int: - if self.flows: - return len(self.flows) - return 0 + if self.current_thread: + current = 1 + else: + current = 0 + return current + len(self.flows) @command.command("replay.client.stop") def stop_replay(self) -> None: @@ -32,7 +33,7 @@ class ClientPlayback: """ Replay requests from flows. """ - self.flows = flows + self.flows = list(flows) ctx.master.addons.trigger("update", []) @command.command("replay.client.file") @@ -54,13 +55,16 @@ class ClientPlayback: self.start_replay(flows) def tick(self): - if self.current_thread and not self.current_thread.is_alive(): + current_is_done = self.current_thread and not self.current_thread.is_alive() + can_start_new = not self.current_thread or current_is_done + will_start_new = can_start_new and self.flows + + if current_is_done: self.current_thread = None - if self.flows and not self.current_thread: + ctx.master.addons.trigger("update", []) + if will_start_new: f = self.flows.pop(0) self.current_thread = ctx.master.replay_request(f) ctx.master.addons.trigger("update", [f]) - self.has_replayed = True - if self.has_replayed: - if not self.flows and not self.current_thread: - ctx.master.addons.trigger("processing_complete") + if current_is_done and not will_start_new: + ctx.master.addons.trigger("processing_complete") diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py index fd673a6f..00bb31c9 100644 --- a/mitmproxy/proxy/protocol/http_replay.py +++ b/mitmproxy/proxy/protocol/http_replay.py @@ -42,6 +42,7 @@ class RequestReplayThread(basethread.BaseThread): super().__init__( "RequestReplay (%s)" % f.request.url ) + self.daemon = True def run(self): r = self.f.request diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py index 6089b2d5..2dc7eb92 100644 --- a/test/mitmproxy/addons/test_clientplayback.py +++ b/test/mitmproxy/addons/test_clientplayback.py @@ -36,9 +36,12 @@ class TestClientPlayback: assert rp.called assert cp.current_thread - cp.flows = None - cp.current_thread = None + cp.flows = [] + cp.current_thread.is_alive.return_value = False + assert cp.count() == 1 cp.tick() + assert cp.count() == 0 + assert tctx.master.has_event("update") assert tctx.master.has_event("processing_complete") cp.current_thread = MockThread() |