Make the http server handle coroutine-making REST servlets (#5475)

2 files changed, 47 insertions, 18 deletions

diff --git a/tests/rest/media/v1/test_media_storage.py b/tests/rest/media/v1/test_media_storage.py
index e2d418b1df..39c9342423 100644
--- a/tests/rest/media/v1/test_media_storage.py
+++ b/tests/rest/media/v1/test_media_storage.py
@@ -22,7 +22,6 @@ from binascii import unhexlify
 from mock import Mock
 from six.moves.urllib import parse
 
-from twisted.internet import defer, reactor
 from twisted.internet.defer import Deferred
 
 from synapse.rest.media.v1._base import FileInfo
@@ -34,15 +33,17 @@ from synapse.util.logcontext import make_deferred_yieldable
 from tests import unittest
 
 
-class MediaStorageTests(unittest.TestCase):
-    def setUp(self):
+class MediaStorageTests(unittest.HomeserverTestCase):
+
+    needs_threadpool = True
+
+    def prepare(self, reactor, clock, hs):
         self.test_dir = tempfile.mkdtemp(prefix="synapse-tests-")
+        self.addCleanup(shutil.rmtree, self.test_dir)
 
         self.primary_base_path = os.path.join(self.test_dir, "primary")
         self.secondary_base_path = os.path.join(self.test_dir, "secondary")
 
-        hs = Mock()
-        hs.get_reactor = Mock(return_value=reactor)
         hs.config.media_store_path = self.primary_base_path
 
         storage_providers = [FileStorageProviderBackend(hs, self.secondary_base_path)]
@@ -52,10 +53,6 @@ class MediaStorageTests(unittest.TestCase):
             hs, self.primary_base_path, self.filepaths, storage_providers
         )
 
-    def tearDown(self):
-        shutil.rmtree(self.test_dir)
-
-    @defer.inlineCallbacks
     def test_ensure_media_is_in_local_cache(self):
         media_id = "some_media_id"
         test_body = "Test\n"
@@ -73,7 +70,15 @@ class MediaStorageTests(unittest.TestCase):
         # Now we run ensure_media_is_in_local_cache, which should copy the file
         # to the local cache.
         file_info = FileInfo(None, media_id)
-        local_path = yield self.media_storage.ensure_media_is_in_local_cache(file_info)
+
+        # This uses a real blocking threadpool so we have to wait for it to be
+        # actually done :/
+        x = self.media_storage.ensure_media_is_in_local_cache(file_info)
+
+        # Hotloop until the threadpool does its job...
+        self.wait_on_thread(x)
+
+        local_path = self.get_success(x)
 
         self.assertTrue(os.path.exists(local_path))
 
diff --git a/tests/unittest.py b/tests/unittest.py
index 36df43c137..d26804b5b5 100644
--- a/tests/unittest.py
+++ b/tests/unittest.py
@@ -17,6 +17,7 @@ import gc
 import hashlib
 import hmac
 import logging
+import time
 
 from mock import Mock
 
@@ -24,7 +25,8 @@ from canonicaljson import json
 
 import twisted
 import twisted.logger
-from twisted.internet.defer import Deferred
+from twisted.internet.defer import Deferred, succeed
+from twisted.python.threadpool import ThreadPool
 from twisted.trial import unittest
 
 from synapse.api.constants import EventTypes
@@ -164,6 +166,7 @@ class HomeserverTestCase(TestCase):
 
     servlets = []
     hijack_auth = True
+    needs_threadpool = False
 
     def setUp(self):
         """
@@ -192,15 +195,19 @@ class HomeserverTestCase(TestCase):
             if self.hijack_auth:
 
                 def get_user_by_access_token(token=None, allow_guest=False):
-                    return {
-                        "user": UserID.from_string(self.helper.auth_user_id),
-                        "token_id": 1,
-                        "is_guest": False,
-                    }
+                    return succeed(
+                        {
+                            "user": UserID.from_string(self.helper.auth_user_id),
+                            "token_id": 1,
+                            "is_guest": False,
+                        }
+                    )
 
                 def get_user_by_req(request, allow_guest=False, rights="access"):
-                    return create_requester(
-                        UserID.from_string(self.helper.auth_user_id), 1, False, None
+                    return succeed(
+                        create_requester(
+                            UserID.from_string(self.helper.auth_user_id), 1, False, None
+                        )
                     )
 
                 self.hs.get_auth().get_user_by_req = get_user_by_req
@@ -209,9 +216,26 @@ class HomeserverTestCase(TestCase):
                     return_value="1234"
                 )
 
+        if self.needs_threadpool:
+            self.reactor.threadpool = ThreadPool()
+            self.addCleanup(self.reactor.threadpool.stop)
+            self.reactor.threadpool.start()
+
         if hasattr(self, "prepare"):
             self.prepare(self.reactor, self.clock, self.hs)
 
+    def wait_on_thread(self, deferred, timeout=10):
+        """
+        Wait until a Deferred is done, where it's waiting on a real thread.
+        """
+        start_time = time.time()
+
+        while not deferred.called:
+            if start_time + timeout < time.time():
+                raise ValueError("Timed out waiting for threadpool")
+            self.reactor.advance(0.01)
+            time.sleep(0.01)
+
     def make_homeserver(self, reactor, clock):
         """
         Make and return a homeserver.