1 files changed, 29 insertions, 6 deletions
diff --git a/tests/server.py b/tests/server.py
index e249668d21..c63b2c3100 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -24,6 +24,7 @@ class FakeChannel(object):
     """
 
     result = attr.ib(default=attr.Factory(dict))
+    _producer = None
 
     @property
     def json_body(self):
@@ -49,6 +50,15 @@ class FakeChannel(object):
 
         self.result["body"] += content
 
+    def registerProducer(self, producer, streaming):
+        self._producer = producer
+
+    def unregisterProducer(self):
+        if self._producer is None:
+            return
+
+        self._producer = None
+
     def requestDone(self, _self):
         self.result["done"] = True
 
@@ -111,14 +121,19 @@ def make_request(method, path, content=b""):
     return req, channel
 
 
-def wait_until_result(clock, channel, timeout=100):
+def wait_until_result(clock, request, timeout=100):
     """
-    Wait until the channel has a result.
+    Wait until the request is finished.
     """
     clock.run()
     x = 0
 
-    while not channel.result:
+    while not request.finished:
+
+        # If there's a producer, tell it to resume producing so we get content
+        if request._channel._producer:
+            request._channel._producer.resumeProducing()
+
         x += 1
 
         if x > timeout:
@@ -129,13 +144,14 @@ def wait_until_result(clock, channel, timeout=100):
 
 def render(request, resource, clock):
     request.render(resource)
-    wait_until_result(clock, request._channel)
+    wait_until_result(clock, request)
 
 
 class ThreadedMemoryReactorClock(MemoryReactorClock):
     """
     A MemoryReactorClock that supports callFromThread.
     """
+
     def callFromThread(self, callback, *args, **kwargs):
         """
         Make the callback fire in the next reactor iteration.
@@ -146,12 +162,15 @@ class ThreadedMemoryReactorClock(MemoryReactorClock):
         return d
 
 
-def setup_test_homeserver(*args, **kwargs):
+def setup_test_homeserver(cleanup_func, *args, **kwargs):
     """
     Set up a synchronous test server, driven by the reactor used by
     the homeserver.
     """
-    d = _sth(*args, **kwargs).result
+    d = _sth(cleanup_func, *args, **kwargs).result
+
+    if isinstance(d, Failure):
+        d.raiseException()
 
     # Make the thread pool synchronous.
     clock = d.get_clock()
@@ -184,9 +203,13 @@ def setup_test_homeserver(*args, **kwargs):
         """
         Threadless thread pool.
         """
+
         def start(self):
             pass
 
+        def stop(self):
+            pass
+
         def callInThreadWithCallback(self, onResult, function, *args, **kwargs):
             def _(res):
                 if isinstance(res, Failure):