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: [PATCH] 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().
---
 changelog.d/8760.misc            |  1 +
 tests/rest/client/v1/utils.py    | 10 ++++++----
 tests/server.py                  | 11 ++++++++++-
 tests/storage/test_client_ips.py | 13 +++++++------
 4 files changed, 24 insertions(+), 11 deletions(-)
 create mode 100644 changelog.d/8760.misc

diff --git a/changelog.d/8760.misc b/changelog.d/8760.misc
new file mode 100644
index 000000000000..54502e9b9029
--- /dev/null
+++ b/changelog.d/8760.misc
@@ -0,0 +1 @@
+Refactor test utilities for injecting HTTP requests.
diff --git a/tests/rest/client/v1/utils.py b/tests/rest/client/v1/utils.py
index afaf9f7b851c..1b2d0497a607 100644
--- a/tests/rest/client/v1/utils.py
+++ b/tests/rest/client/v1/utils.py
@@ -296,10 +296,12 @@ def upload_media(
         image_length = len(image_data)
         path = "/_matrix/media/r0/upload?filename=%s" % (filename,)
         request, channel = make_request(
-            self.hs.get_reactor(), "POST", path, content=image_data, access_token=tok
-        )
-        request.requestHeaders.addRawHeader(
-            b"Content-Length", str(image_length).encode("UTF-8")
+            self.hs.get_reactor(),
+            "POST",
+            path,
+            content=image_data,
+            access_token=tok,
+            custom_headers=[(b"Content-Length", str(image_length))],
         )
         request.render(resource)
         self.hs.get_reactor().pump([100])
diff --git a/tests/server.py b/tests/server.py
index 3dd2cfc072a3..ef03109a6ce2 100644
--- a/tests/server.py
+++ b/tests/server.py
@@ -2,7 +2,7 @@
 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
diff --git a/tests/storage/test_client_ips.py b/tests/storage/test_client_ips.py
index e96ca1c8cab9..efca43ec78a2 100644
--- a/tests/storage/test_client_ips.py
+++ b/tests/storage/test_client_ips.py
@@ -21,6 +21,7 @@
 from synapse.rest.client.v1 import login
 
 from tests import unittest
+from tests.server import make_request
 from tests.test_utils import make_awaitable
 from tests.unittest import override_config
 
@@ -408,17 +409,17 @@ def _runtest(self, headers, expected_ip, make_request_args):
         # Advance to a known time
         self.reactor.advance(123456 - self.reactor.seconds())
 
-        request, channel = self.make_request(
+        headers1 = {b"User-Agent": b"Mozzila pizza"}
+        headers1.update(headers)
+
+        request, channel = make_request(
+            self.reactor,
             "GET",
             "/_matrix/client/r0/admin/users/" + self.user_id,
             access_token=access_token,
+            custom_headers=headers1.items(),
             **make_request_args,
         )
-        request.requestHeaders.addRawHeader(b"User-Agent", b"Mozzila pizza")
-
-        # Add the optional headers
-        for h, v in headers.items():
-            request.requestHeaders.addRawHeader(h, v)
         self.render(request)
 
         # Advance so the save loop occurs