client replay: expad and consolidate tests

2 files changed, 105 insertions, 76 deletions

diff --git a/test/mitmproxy/addons/test_clientplayback.py b/test/mitmproxy/addons/test_clientplayback.py
index a63bec53..1b385e23 100644
--- a/test/mitmproxy/addons/test_clientplayback.py
+++ b/test/mitmproxy/addons/test_clientplayback.py
@@ -1,13 +1,16 @@
+import time
 import pytest
-from unittest import mock
 
-from mitmproxy.test import tflow
+from mitmproxy.test import tflow, tutils
 from mitmproxy import io
 from mitmproxy import exceptions
+from mitmproxy.net import http as net_http
 
 from mitmproxy.addons import clientplayback
 from mitmproxy.test import taddons
 
+from .. import tservers
+
 
 def tdump(path, flows):
     with open(path, "wb") as f:
@@ -21,77 +24,80 @@ class MockThread():
         return False
 
 
-class TestClientPlayback:
-    # @staticmethod
-    # def wait_until_not_live(flow):
-    #     """
-    #     Race condition: We don't want to replay the flow while it is still live.
-    #     """
-    #     s = time.time()
-    #     while flow.live:
-    #         time.sleep(0.001)
-    #         if time.time() - s > 5:
-    #             raise RuntimeError("Flow is live for too long.")
-
-    # def test_replay(self):
-    #     assert self.pathod("304").status_code == 304
-    #     assert len(self.master.state.flows) == 1
-    #     l = self.master.state.flows[-1]
-    #     assert l.response.status_code == 304
-    #     l.request.path = "/p/305"
-    #     self.wait_until_not_live(l)
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert l.response.status_code == 305
-
-    #     # Disconnect error
-    #     l.request.path = "/p/305:d0"
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert rt
-    #     if isinstance(self, tservers.HTTPUpstreamProxyTest):
-    #         assert l.response.status_code == 502
-    #     else:
-    #         assert l.error
-
-    #     # Port error
-    #     l.request.port = 1
-    #     # In upstream mode, we get a 502 response from the upstream proxy server.
-    #     # In upstream mode with ssl, the replay will fail as we cannot establish
-    #     # SSL with the upstream proxy.
-    #     rt = self.master.replay_request(l, block=True)
-    #     assert rt
-    #     if isinstance(self, tservers.HTTPUpstreamProxyTest):
-    #         assert l.response.status_code == 502
-    #     else:
-    #         assert l.error
-
-    # def test_replay(self):
-    #     opts = options.Options()
-    #     fm = master.Master(opts)
-    #     f = tflow.tflow(resp=True)
-    #     f.request.content = None
-    #     with pytest.raises(ReplayException, match="missing"):
-    #         fm.replay_request(f)
-
-    #     f.request = None
-    #     with pytest.raises(ReplayException, match="request"):
-    #         fm.replay_request(f)
-
-    #     f.intercepted = True
-    #     with pytest.raises(ReplayException, match="intercepted"):
-    #         fm.replay_request(f)
-
-    #     f.live = True
-    #     with pytest.raises(ReplayException, match="live"):
-    #         fm.replay_request(f)
-
-    #     req = tutils.treq(headers=net_http.Headers(((b":authority", b"foo"), (b"header", b"qvalue"), (b"content-length", b"7"))))
-    #     f = tflow.tflow(req=req)
-    #     f.request.http_version = "HTTP/2.0"
-    #     with mock.patch('mitmproxy.proxy.protocol.http_replay.RequestReplayThread.run'):
-    #         rt = fm.replay_request(f)
-    #         assert rt.f.request.http_version == "HTTP/1.1"
-    #         assert ":authority" not in rt.f.request.headers
+class TBase(tservers.HTTPProxyTest):
+    @staticmethod
+    def wait_response(flow):
+        """
+        Race condition: We don't want to replay the flow while it is still live.
+        """
+        s = time.time()
+        while True:
+            if flow.response or flow.error:
+                break
+            time.sleep(0.001)
+            if time.time() - s > 5:
+                raise RuntimeError("Flow is live for too long.")
+
+    @staticmethod
+    def reset(f):
+        f.live = False
+        f.repsonse = False
+        f.error = False
+
+    def addons(self):
+        return [clientplayback.ClientPlayback()]
+
+    def test_replay(self):
+        cr = self.master.addons.get("clientplayback")
+
+        assert self.pathod("304").status_code == 304
+        assert len(self.master.state.flows) == 1
+        l = self.master.state.flows[-1]
+        assert l.response.status_code == 304
+        l.request.path = "/p/305"
+        cr.start_replay([l])
+        self.wait_response(l)
+        assert l.response.status_code == 305
+
+        # Disconnect error
+        cr.stop_replay()
+        self.reset(l)
+        l.request.path = "/p/305:d0"
+        cr.start_replay([l])
+        self.wait_response(l)
+        if isinstance(self, tservers.HTTPUpstreamProxyTest):
+            assert l.response.status_code == 502
+        else:
+            assert l.error
+
+        # # Port error
+        cr.stop_replay()
+        self.reset(l)
+        l.request.port = 1
+        # In upstream mode, we get a 502 response from the upstream proxy server.
+        # In upstream mode with ssl, the replay will fail as we cannot establish
+        # SSL with the upstream proxy.
+        cr.start_replay([l])
+        self.wait_response(l)
+        if isinstance(self, tservers.HTTPUpstreamProxyTest):
+            assert l.response.status_code == 502
+        else:
+            assert l.error
+
+
+class TestHTTPProxy(TBase, tservers.HTTPProxyTest):
+    pass
+
+
+class TestHTTPSProxy(TBase, tservers.HTTPProxyTest):
+    ssl = True
+
+
+class TestUpstreamProxy(TBase, tservers.HTTPUpstreamProxyTest):
+    pass
 
+
+class TestClientPlayback:
     def test_load_file(self, tmpdir):
         cp = clientplayback.ClientPlayback()
         with taddons.context(cp):
@@ -133,13 +139,37 @@ class TestClientPlayback:
             f.request.raw_content = None
             assert "missing content" in cp.check(f)
 
-    def test_playback(self):
+    @pytest.mark.asyncio
+    async def test_playback(self):
         cp = clientplayback.ClientPlayback()
-        with taddons.context(cp):
+        with taddons.context(cp) as ctx:
             assert cp.count() == 0
             f = tflow.tflow(resp=True)
             cp.start_replay([f])
             assert cp.count() == 1
 
             cp.stop_replay()
-            assert cp.count() == 0
\ No newline at end of file
+            assert cp.count() == 0
+
+            f.live = True
+            cp.start_replay([f])
+            assert cp.count() == 0
+            await ctx.master.await_log("live")
+
+    def test_http2(self):
+        cp = clientplayback.ClientPlayback()
+        with taddons.context(cp):
+            req = tutils.treq(
+                headers = net_http.Headers(
+                    (
+                        (b":authority", b"foo"),
+                        (b"header", b"qvalue"),
+                        (b"content-length", b"7")
+                    )
+                )
+            )
+            f = tflow.tflow(req=req)
+            f.request.http_version = "HTTP/2.0"
+            cp.start_replay([f])
+            assert f.request.http_version == "HTTP/1.1"
+            assert ":authority" not in f.request.headers
diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py
index b272861c..3d18987d 100644
--- a/test/mitmproxy/tools/web/test_app.py
+++ b/test/mitmproxy/tools/web/test_app.py
@@ -9,7 +9,6 @@ import tornado.testing
 from tornado import httpclient
 from tornado import websocket
 
-from mitmproxy import exceptions
 from mitmproxy import options
 from mitmproxy.test import tflow
 from mitmproxy.tools.web import app