From ebc405446e6615d6187a2e29cb33f27dd5bd0841 Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
Date: Mon, 16 Nov 2020 14:45:22 +0000
Subject: Add a `custom_headers` param to `make_request` (#8760)

Some tests want to set some custom HTTP request headers, so provide a way to do
that before calling requestReceived().
---
 tests/server.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

(limited to 'tests/server.py')

diff --git a/tests/server.py b/tests/server.py
index 3dd2cfc072..ef03109a6c 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -2,7 +2,7 @@ import json
 import logging
 from collections import deque
 from io import SEEK_END, BytesIO
-from typing import Callable
+from typing import Callable, Iterable, Optional, Tuple, Union
 
 import attr
 from typing_extensions import Deque
@@ -139,6 +139,9 @@ def make_request(
     shorthand=True,
     federation_auth_origin=None,
     content_is_form=False,
+    custom_headers: Optional[
+        Iterable[Tuple[Union[bytes, str], Union[bytes, str]]]
+    ] = None,
 ):
     """
     Make a web request using the given method and path, feed it the
@@ -157,6 +160,8 @@ def make_request(
         content_is_form: Whether the content is URL encoded form data. Adds the
             'Content-Type': 'application/x-www-form-urlencoded' header.
 
+        custom_headers: (name, value) pairs to add as request headers
+
     Returns:
         Tuple[synapse.http.site.SynapseRequest, channel]
     """
@@ -211,6 +216,10 @@ def make_request(
             # Assume the body is JSON
             req.requestHeaders.addRawHeader(b"Content-Type", b"application/json")
 
+    if custom_headers:
+        for k, v in custom_headers:
+            req.requestHeaders.addRawHeader(k, v)
+
     req.requestReceived(method, path, b"1.1")
 
     return req, channel
-- 
cgit 1.5.1


From f125895475aeee9447f3988ecbd8bfd1836545bf Mon Sep 17 00:00:00 2001
From: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
Date: Mon, 16 Nov 2020 18:21:47 +0000
Subject: Move `wait_until_result` into `FakeChannel` (#8758)

FakeChannel has everything we need, and this more accurately models the real
flow.
---
 changelog.d/8758.misc                         |  1 +
 tests/rest/key/v2/test_remote_key_resource.py |  6 ++--
 tests/server.py                               | 42 +++++++++++++--------------
 3 files changed, 24 insertions(+), 25 deletions(-)
 create mode 100644 changelog.d/8758.misc

(limited to 'tests/server.py')

diff --git a/changelog.d/8758.misc b/changelog.d/8758.misc
new file mode 100644
index 0000000000..54502e9b90
--- /dev/null
+++ b/changelog.d/8758.misc
@@ -0,0 +1 @@
+Refactor test utilities for injecting HTTP requests.
diff --git a/tests/rest/key/v2/test_remote_key_resource.py b/tests/rest/key/v2/test_remote_key_resource.py
index 6671cbd32d..fbcf8d5b86 100644
--- a/tests/rest/key/v2/test_remote_key_resource.py
+++ b/tests/rest/key/v2/test_remote_key_resource.py
@@ -32,7 +32,7 @@ from synapse.util.httpresourcetree import create_resource_tree
 from synapse.util.stringutils import random_string
 
 from tests import unittest
-from tests.server import FakeChannel, wait_until_result
+from tests.server import FakeChannel
 from tests.utils import default_config
 
 
@@ -94,7 +94,7 @@ class RemoteKeyResourceTestCase(BaseRemoteKeyResourceTestCase):
             % (server_name.encode("utf-8"), key_id.encode("utf-8")),
             b"1.1",
         )
-        wait_until_result(self.reactor, req)
+        channel.await_result()
         self.assertEqual(channel.code, 200)
         resp = channel.json_body
         return resp
@@ -190,7 +190,7 @@ class EndToEndPerspectivesTests(BaseRemoteKeyResourceTestCase):
             req.requestReceived(
                 b"POST", path.encode("utf-8"), b"1.1",
             )
-            wait_until_result(self.reactor, req)
+            channel.await_result()
             self.assertEqual(channel.code, 200)
             resp = channel.json_body
             return resp
diff --git a/tests/server.py b/tests/server.py
index ef03109a6c..18cb8b2d72 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -117,6 +117,25 @@ class FakeChannel:
     def transport(self):
         return self
 
+    def await_result(self, timeout: int = 100) -> None:
+        """
+        Wait until the request is finished.
+        """
+        self._reactor.run()
+        x = 0
+
+        while not self.result.get("done"):
+            # If there's a producer, tell it to resume producing so we get content
+            if self._producer:
+                self._producer.resumeProducing()
+
+            x += 1
+
+            if x > timeout:
+                raise TimedOutException("Timed out waiting for request to finish.")
+
+            self._reactor.advance(0.1)
+
 
 class FakeSite:
     """
@@ -225,30 +244,9 @@ def make_request(
     return req, channel
 
 
-def wait_until_result(clock, request, timeout=100):
-    """
-    Wait until the request is finished.
-    """
-    clock.run()
-    x = 0
-
-    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:
-            raise TimedOutException("Timed out waiting for request to finish.")
-
-        clock.advance(0.1)
-
-
 def render(request, resource, clock):
     request.render(resource)
-    wait_until_result(clock, request)
+    request._channel.await_result()
 
 
 @implementer(IReactorPluggableNameResolver)
-- 
cgit 1.5.1