From 28a8ddc0e827149c76a738e6363713ed2ebc0eac Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@corte.si>
Date: Mon, 2 Apr 2018 09:46:02 +1200
Subject: [PATCH] asyncio: fix a few remaining issues in proxy/test_server.py

---
 test/mitmproxy/proxy/test_server.py |  2 +-
 test/mitmproxy/tservers.py          | 21 +++++++++++++--------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/test/mitmproxy/proxy/test_server.py b/test/mitmproxy/proxy/test_server.py
index 9e2762946..aed4a7745 100644
--- a/test/mitmproxy/proxy/test_server.py
+++ b/test/mitmproxy/proxy/test_server.py
@@ -987,7 +987,7 @@ class TestUpstreamProxySSL(
         Client <- HTTPS -> Proxy <- HTTP -> Proxy <- HTTPS -> Server
         """
         self.set_addons(RewriteToHttp())
-        self.set_addons(RewriteToHttps())
+        self.chain[1].set_addons(RewriteToHttps())
         p = self.pathoc()
         with p.connect():
             resp = p.request("get:'/p/418'")
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index f3270e5c1..18e4f27d3 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -117,6 +117,11 @@ class ProxyThread(threading.Thread):
         )
         self.tmaster.run()
 
+    def set_addons(self, *addons):
+        self.tmaster.reset(addons)
+        self.tmaster.addons.trigger("tick")
+
+
 
 class ProxyTestBase:
     # Test Configuration
@@ -180,8 +185,7 @@ class ProxyTestBase:
         )
 
     def set_addons(self, *addons):
-        self.proxy.tmaster.reset(addons)
-        self.proxy.tmaster.addons.trigger("tick")
+        self.proxy.set_addons(*addons)
 
     def addons(self):
         """
@@ -337,8 +341,7 @@ class SocksModeTest(HTTPProxyTest):
         return opts
 
 
-class ChainProxyTest(ProxyTestBase):
-
+class HTTPUpstreamProxyTest(HTTPProxyTest):
     """
     Chain three instances of mitmproxy in a row to test upstream mode.
     Proxy order is cls.proxy -> cls.chain[0] -> cls.chain[1]
@@ -357,6 +360,12 @@ class ChainProxyTest(ProxyTestBase):
             proxy = ProxyThread(cls.masterclass, opts)
             proxy.start()
             cls.chain.insert(0, proxy)
+            while 1:
+                if(
+                    proxy.event_loop and
+                    proxy.event_loop.is_running()
+                ):
+                    break
 
         super().setup_class()
 
@@ -380,7 +389,3 @@ class ChainProxyTest(ProxyTestBase):
                 mode="upstream:" + s,
             )
         return opts
-
-
-class HTTPUpstreamProxyTest(ChainProxyTest, HTTPProxyTest):
-    pass