From 8df82f50e48f0716cef72065128c0f1d521ba19a Mon Sep 17 00:00:00 2001 From: Alex Vandiver Date: Thu, 6 May 2021 16:57:46 -0700 Subject: [PATCH] outgoing_http: Provide a convenient way to set default headers. --- zerver/lib/outgoing_http.py | 4 +++- zerver/lib/outgoing_webhook.py | 6 +----- zerver/tests/test_outgoing_http.py | 22 +++++++++++++++++++++- 3 files changed, 25 insertions(+), 7 deletions(-) diff --git a/zerver/lib/outgoing_http.py b/zerver/lib/outgoing_http.py index b2b4f7abdb..a5e4645f60 100644 --- a/zerver/lib/outgoing_http.py +++ b/zerver/lib/outgoing_http.py @@ -5,11 +5,13 @@ from urllib3 import HTTPResponse class OutgoingSession(requests.Session): - def __init__(self, role: str, timeout: int) -> None: + def __init__(self, role: str, timeout: int, headers: Optional[Dict[str, str]] = None) -> None: super().__init__() outgoing_adapter = OutgoingHTTPAdapter(role=role, timeout=timeout) self.mount("http://", outgoing_adapter) self.mount("https://", outgoing_adapter) + if headers: + self.headers.update(headers) class OutgoingHTTPAdapter(requests.adapters.HTTPAdapter): diff --git a/zerver/lib/outgoing_webhook.py b/zerver/lib/outgoing_webhook.py index 0275183085..b571575cec 100644 --- a/zerver/lib/outgoing_webhook.py +++ b/zerver/lib/outgoing_webhook.py @@ -36,11 +36,7 @@ class OutgoingWebhookServiceInterface(metaclass=abc.ABCMeta): self.session: requests.Session = OutgoingSession( role="webhook", timeout=10, - ) - self.session.headers.update( - { - "User-Agent": "ZulipOutgoingWebhook/" + ZULIP_VERSION, - } + headers={"User-Agent": "ZulipOutgoingWebhook/" + ZULIP_VERSION}, ) @abc.abstractmethod diff --git a/zerver/tests/test_outgoing_http.py b/zerver/tests/test_outgoing_http.py index dcd13ba561..aec23f880c 100644 --- a/zerver/tests/test_outgoing_http.py +++ b/zerver/tests/test_outgoing_http.py @@ -48,15 +48,35 @@ class RequestMockWithTimeoutAsHeader(responses.RequestsMock): class TestOutgoingHttp(ZulipTestCase): + def test_headers(self) -> None: + with RequestMockWithProxySupport() as mock_requests: + mock_requests.add(responses.GET, "http://example.com/") + OutgoingSession(role="testing", timeout=1, headers={"X-Foo": "bar"}).get( + "http://example.com/" + ) + self.assertEqual(len(mock_requests.calls), 1) + headers = mock_requests.calls[0].request.headers + # We don't see a proxy header with no proxy set + self.assertFalse("X-Smokescreen-Role" in headers) + self.assertEqual(headers["X-Foo"], "bar") + @mock.patch.dict(os.environ, {"http_proxy": "http://localhost:4242"}) def test_proxy_headers(self) -> None: with RequestMockWithProxySupport() as mock_requests: mock_requests.add(responses.GET, "http://localhost:4242/") - OutgoingSession(role="testing", timeout=1).get("http://example.com/") + OutgoingSession(role="testing", timeout=1, headers={"X-Foo": "bar"}).get( + "http://example.com/" + ) self.assertEqual(len(mock_requests.calls), 1) headers = mock_requests.calls[0].request.headers self.assertEqual(headers["X-Smokescreen-Role"], "testing") + # We don't see the request-level headers in the proxy + # request. This isn't a _true_ test because we're + # fiddling the headers above, instead of urllib3 actually + # setting them. + self.assertFalse("X-Foo" in headers) + def test_timeouts(self) -> None: with RequestMockWithTimeoutAsHeader() as mock_requests: mock_requests.add(responses.GET, "http://example.com/")