| 181 | |
| 182 | |
| 183 | class _TestHandler(CopilotRequestHandler): |
| 184 | def __init__(self, upstream: _Upstream, counters: _Counters) -> None: |
| 185 | self._upstream = upstream |
| 186 | self._counters = counters |
| 187 | self._client = httpx.AsyncClient(timeout=None, follow_redirects=False) |
| 188 | |
| 189 | def _rewrite_http(self, url: httpx.URL) -> httpx.URL: |
| 190 | up = httpx.URL(self._upstream.http_url) |
| 191 | return url.copy_with(scheme=up.scheme, host=up.host, port=up.port) |
| 192 | |
| 193 | def _rewrite_ws(self, url: str) -> str: |
| 194 | parsed = httpx.URL(url) |
| 195 | up = httpx.URL(self._upstream.ws_url) |
| 196 | return str(parsed.copy_with(scheme=up.scheme, host=up.host, port=up.port)) |
| 197 | |
| 198 | async def send_request( |
| 199 | self, request: httpx.Request, ctx: CopilotRequestContext |
| 200 | ) -> httpx.Response: |
| 201 | self._counters.http_requests += 1 |
| 202 | headers = dict(request.headers) |
| 203 | headers["x-test-mutated"] = "1" |
| 204 | rewritten = httpx.Request( |
| 205 | request.method, |
| 206 | self._rewrite_http(request.url), |
| 207 | headers=headers, |
| 208 | content=request.content, |
| 209 | ) |
| 210 | response = await self._client.send(rewritten, stream=True) |
| 211 | self._counters.http_responses += 1 |
| 212 | response.headers["x-test-response-mutated"] = "1" |
| 213 | return response |
| 214 | |
| 215 | async def open_websocket(self, ctx: CopilotRequestContext): |
| 216 | ctx.url = self._rewrite_ws(ctx.url) |
| 217 | return _CountingSocketHandler(ctx, self._counters) |
| 218 | |
| 219 | async def aclose(self) -> None: |
| 220 | await self._client.aclose() |
| 221 | |
| 222 | |
| 223 | @dataclass |
no outgoing calls
searching dependent graphs…