Refine semantics of replay_request method.

2 files changed, 6 insertions, 8 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 53644fb1..aa02b5c5 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -1361,7 +1361,7 @@ class FlowMaster(controller.Master):
         if self.stickycookie_state:
             self.stickycookie_state.handle_response(f)
 
-    def replay_request(self, f):
+    def replay_request(self, f, block=False):
         """
             Returns None if successful, or error message if not.
         """
@@ -1382,7 +1382,8 @@ class FlowMaster(controller.Master):
                     self.masterq,
                 )
             rt.start() # pragma: no cover
-        return rt
+            if block:
+                rt.join()
 
     def run_script_hook(self, name, *args, **kwargs):
         if self.script and not self.pause_scripts:
diff --git a/test/test_server.py b/test/test_server.py
index 72be976f..86dc0a4a 100644
--- a/test/test_server.py
+++ b/test/test_server.py
@@ -26,20 +26,17 @@ class SanityMixin:
         l = self.master.state.view[0]
         assert l.response.code == 304
         l.request.path = "/p/305"
-        rt = self.master.replay_request(l)
-        rt.join()
+        rt = self.master.replay_request(l, block=True)
         assert l.response.code == 305
 
         # Disconnect error
         l.request.path = "/p/305:d0"
-        rt = self.master.replay_request(l)
-        rt.join()
+        rt = self.master.replay_request(l, block=True)
         assert l.error
 
         # Port error
         l.request.port = 1
-        rt = self.master.replay_request(l)
-        rt.join()
+        rt = self.master.replay_request(l, block=True)
         assert l.error