fix: port middleware/test_proxy.py to async context

Author: gi0baro

src/sentry/testutils/asserts.py

@@ -2,7 +2,6 @@
 
 from django.db import models
 from django.db.models.functions import Cast
-from django.http import StreamingHttpResponse
 
 from sentry.constants import ObjectStatus
 from sentry.integrations.types import EventLifecycleOutcome
@@ -44,10 +43,7 @@ def assert_commit_shape(commit):
 def assert_status_code(response, minimum: int, maximum: int | None = None):
     # Omit max to assert status_code == minimum.
     maximum = maximum or minimum + 1
-    assert minimum <= response.status_code < maximum, (
-        response.status_code,
-        response.getvalue() if isinstance(response, StreamingHttpResponse) else response.content,
-    )
+    assert minimum <= response.status_code < maximum, response
 
 
 def assert_existing_projects_status(

src/sentry/testutils/cases.py

@@ -17,6 +17,7 @@
 from uuid import UUID, uuid4
 from zlib import compress
 
+import httpx
 import pytest
 import requests
 import responses
@@ -36,7 +37,6 @@
 from django.utils import timezone
 from django.utils.functional import cached_property
 from google.protobuf.timestamp_pb2 import Timestamp
-from requests.utils import CaseInsensitiveDict, get_encoding_from_headers
 from rest_framework import status
 from rest_framework.request import Request
 from rest_framework.response import Response
@@ -678,36 +678,45 @@ def get_cursor_headers(self, response):
     def api_gateway_proxy_stubbed(self):
         """Mocks a fake api gateway proxy that redirects via Client objects"""
 
-        def proxy_raw_request(
-            method: str,
-            url: str,
-            headers: Mapping[str, str],
-            params: Mapping[str, str] | None,
-            data: Any,
-            **kwds: Any,
-        ) -> requests.Response:
-            from django.test.client import Client
-
-            client = Client()
-            extra: Mapping[str, Any] = {
-                f"HTTP_{k.replace('-', '_').upper()}": v for k, v in headers.items()
-            }
-            if params:
-                url += "?" + urlencode(params)
-            with assume_test_silo_mode(SiloMode.CELL):
-                resp = getattr(client, method.lower())(
-                    url, b"".join(data), headers["Content-Type"], **extra
-                )
-            response = requests.Response()
-            response.status_code = resp.status_code
-            response.headers = CaseInsensitiveDict(resp.headers)
-            response.encoding = get_encoding_from_headers(response.headers)
-            response.raw = BytesIO(resp.content)
-            return response
+        from asgiref.sync import sync_to_async
+        from django.test.client import Client
+
+        class MockedProxy:
+            def __init__(self):
+                self.client = Client()
+
+            @staticmethod
+            async def _consume_body(content):
+                ret = b""
+                async for chunk in content:
+                    ret += chunk
+                return ret
+
+            def build_request(self, method, url, headers, params, content, timeout):
+                assert not params
+                target = getattr(self.client, method.lower())
+                content_type = headers.pop("Content-Type", "application/octet-stream")
+                extra: Mapping[str, Any] = {
+                    f"HTTP_{k.replace('-', '_').upper()}": v for k, v in headers.items()
+                }
+                return target, (url, content, content_type), extra
+
+            async def send(self, req, stream, follow_redirects):
+                with assume_test_silo_mode(SiloMode.CELL):
+                    url, content, content_type = req[1]
+                    content = await self._consume_body(content)
+                    resp = await sync_to_async(req[0])(url, content, content_type, **req[2])
+                    wresp = httpx.Response(
+                        status_code=resp.status_code,
+                        headers=dict(resp.headers),
+                        content=resp.content,
+                    )
+                    return wresp
 
+        mock_client = MockedProxy()
         with mock.patch(
-            "sentry.hybridcloud.apigateway.proxy.external_request",
-            new=proxy_raw_request,
+            "sentry.hybridcloud.apigateway.proxy.proxy_client",
+            new=mock_client,
         ):
             yield
 

tests/sentry/middleware/test_proxy.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+import asyncio
 from functools import cached_property
+from unittest.mock import patch
 
 from django.http import HttpRequest
 
 from sentry.middleware.proxy import SetRemoteAddrFromForwardedFor
 from sentry.models.team import Team
 from sentry.silo.base import SiloMode
 from sentry.testutils.cases import APITestCase, TestCase
+from sentry.testutils.helpers.response import close_streaming_response
 from sentry.testutils.silo import assume_test_silo_mode, control_silo_test
 from sentry.types.cell import Cell, RegionCategory
 from sentry.utils import json
@@ -48,6 +51,29 @@ class FakedAPIProxyTest(APITestCase):
     endpoint = "sentry-api-0-organization-teams"
     method = "post"
 
+    def setUp(self) -> None:
+        super().setUp()
+
+        from sentry.hybridcloud.apigateway.middleware import ApiGatewayMiddleware
+
+        _original_middleware = ApiGatewayMiddleware._process_view_inner
+
+        def _process_view_match(self, request, view_func, view_args, view_kwargs):
+            try:
+                asyncio.get_running_loop()
+                return self._process_view_inner(request, view_func, view_args, view_kwargs)
+            except RuntimeError:
+                return self._process_view_sync(request, view_func, view_args, view_kwargs)
+
+        self._middleware_patch = patch.object(
+            ApiGatewayMiddleware, "_process_view_match", _process_view_match
+        )
+        self._middleware_patch.start()
+
+    def tearDown(self) -> None:
+        self._middleware_patch.stop()
+        super().tearDown()
+
     def test_through_api_gateway(self) -> None:
         if SiloMode.get_current_mode() == SiloMode.MONOLITH:
             return
@@ -62,7 +88,7 @@ def test_through_api_gateway(self) -> None:
                 status_code=201,
             )
 
-        result = json.loads(resp.getvalue())
+        result = json.loads(close_streaming_response(resp))
         with assume_test_silo_mode(SiloMode.CELL):
             team = Team.objects.get(id=result["id"])
             assert team.idp_provisioned