from __future__ import annotations import asyncio import base64 import gzip import json import sys import threading import time from collections.abc import Awaitable, Callable, Coroutine, Iterator from typing import TYPE_CHECKING, Any from urllib.parse import parse_qs from uvicorn.server import Server from yarl import URL from tests.unit.server_endpoints import ( BASE_INDEX, GENERIC_RESPONSE, HELLO_WORLD, INCAPSULA, INFINITE_SCROLL, NON_HREF_LINKS, PROBLEMATIC_LINKS, RESOURCE_LOADING_PAGE, ROBOTS_TXT, SECONDARY_INDEX, START_ENQUEUE, START_ENQUEUE_NON_HREF, ) if TYPE_CHECKING: from socket import socket Receive = Callable[[], Awaitable[dict[str, Any]]] Send = Callable[[dict[str, Any]], Coroutine[None, None, None]] PathHandler = Callable[[dict[str, Any], Receive, Send], Coroutine[None, None, None]] def get_headers_dict(scope: dict[str, Any]) -> dict[str, str]: """Extract request headers and return them as a dictionary.""" headers = {} for name, value in scope.get('headers', []): headers[name.decode()] = value.decode() return headers def get_query_params(query_string: bytes) -> dict[str, str]: """Extract and parse query parameters from the request.""" args = parse_qs(query_string.decode(), keep_blank_values=True) result_args = {} for key, values in args.items(): if values: result_args[key] = values[0] return result_args def get_cookies_from_headers(headers: dict[str, Any]) -> dict[str, str]: """Extract cookies from request headers.""" cookies = {} cookie_header: str = headers.get('cookie', '') if cookie_header: for cookie in cookie_header.split(';'): name, value = cookie.strip().split('=') cookies[name] = value return cookies async def send_json_response(send: Send, data: Any, status: int = 200) -> None: """Send a JSON response to the client.""" await send( { 'type': 'http.response.start', 'status': status, 'headers': [[b'content-type', b'application/json']], } ) await send({'type': 'http.response.body', 'body': json.dumps(data, indent=2).encode()}) async def send_html_response(send: Send, html_content: bytes, status: int = 200) -> None: """Send an HTML response to the client.""" await send( { 'type': 'http.response.start', 'status': status, 'headers': [[b'content-type', b'text/html; charset=utf-8']], } ) await send({'type': 'http.response.body', 'body': html_content}) async def app(scope: dict[str, Any], receive: Receive, send: Send) -> None: """Main ASGI application handler that routes requests to specific handlers. Args: scope: The ASGI connection scope. receive: The ASGI receive function. send: The ASGI send function. """ assert scope['type'] == 'http' paths: dict[str, PathHandler] = { 'start_enqueue': start_enqueue_endpoint, 'start_enqueue_non_href': start_enqueue_non_href_endpoint, 'sub_index': secondary_index_endpoint, 'incapsula': incapsula_endpoint, 'page_1': generic_response_endpoint, 'page_2': generic_response_endpoint, 'page_3': generic_response_endpoint, 'base_page': base_index_endpoint, 'problematic_links': problematic_links_endpoint, 'non_href_links': non_href_links_endpoint, 'set_cookies': set_cookies, 'set_complex_cookies': set_complex_cookies, 'cookies': get_cookies, 'status': echo_status, 'headers': echo_headers, 'user-agent': echo_user_agent, 'echo_content': echo_content, 'sitemap.txt': echo_content, 'sitemap.xml': echo_content, 'sitemap.xml.gz': echo_content, 'get': get_echo, 'post': post_echo, 'redirect': redirect_to_url, 'json': hello_world_json, 'xml': hello_world_xml, 'robots.txt': robots_txt, 'get_compressed': get_compressed, 'slow': slow_response, 'infinite_scroll': infinite_scroll_endpoint, 'resource_loading_page': resource_loading_endpoint, } path = URL(scope['path']).parts[1] # Route requests to appropriate handlers if path in paths: path_func = paths[path] await path_func(scope, receive, send) else: await hello_world(scope, receive, send) async def get_cookies(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests to retrieve cookies sent in the request.""" headers = get_headers_dict(scope) cookies = get_cookies_from_headers(headers) await send_json_response(send, {'cookies': cookies}) async def set_cookies(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests to set cookies from query parameters and redirect.""" query_params = get_query_params(scope.get('query_string', b'')) headers = [ [b'content-type', b'text/plain; charset=utf-8'], [b'location', b'/cookies'], # Redirect header ] for key, values in query_params.items(): if values: # Only add if there's at least one value cookie_value = f'{key}={values[0]}; Path=/' headers.append([b'set-cookie', cookie_value.encode()]) await send( { 'type': 'http.response.start', 'status': 302, # 302 Found for redirect 'headers': headers, } ) await send({'type': 'http.response.body', 'body': b'Redirecting to get_cookies...'}) async def hello_world(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle basic requests with a simple HTML response.""" await send_html_response( send, HELLO_WORLD, ) async def hello_world_json(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle basic requests with a simple JSON response.""" await send_json_response( send, {'hello': 'world'}, ) async def hello_world_xml(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle basic requests with a simple XML response.""" await send_html_response( send, b""" world""", ) async def post_echo(scope: dict[str, Any], receive: Receive, send: Send) -> None: """Echo back POST request details similar to httpbin.org/post.""" # Extract basic request info path = scope.get('path', '') query_string = scope.get('query_string', b'') args = get_query_params(query_string) # Extract headers and cookies headers = get_headers_dict(scope) # Read the request body body = b'' form = {} json_data = None more_body = True while more_body: message = await receive() if message['type'] == 'http.request': body += message.get('body', b'') more_body = message.get('more_body', False) # Parse body based on content type content_type = headers.get('content-type', '').lower() if body and 'application/json' in content_type: json_data = json.loads(body.decode()) if body and 'application/x-www-form-urlencoded' in content_type: form_data = parse_qs(body.decode()) for key, values in form_data.items(): form[key] = values[0] if len(values) == 1 else values body_text = '' if form else body.decode('utf-8', errors='replace') # Prepare response response = { 'args': args, 'data': body_text, 'files': {}, # Not handling multipart file uploads 'form': form, 'headers': headers, 'json': json_data, 'origin': headers.get('host', ''), 'url': f'http://{headers["host"]}{path}', } await send_json_response(send, response) async def echo_status(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Echo the status code from the URL path.""" status_code = int(scope['path'].replace('/status/', '')) await send( { 'type': 'http.response.start', 'status': status_code, 'headers': [[b'content-type', b'text/plain']], } ) await send({'type': 'http.response.body', 'body': b''}) async def echo_headers(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Echo back the request headers as JSON.""" headers = get_headers_dict(scope) await send_json_response(send, headers) async def start_enqueue_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the main page with links.""" await send_html_response( send, START_ENQUEUE, ) async def secondary_index_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the secondary page with links.""" await send_html_response( send, SECONDARY_INDEX, ) async def incapsula_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for a page with an incapsula iframe.""" await send_html_response( send, INCAPSULA, ) async def generic_response_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests with a generic HTML response.""" await send_html_response( send, GENERIC_RESPONSE, ) async def problematic_links_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests with a page containing problematic links.""" await send_html_response( send, PROBLEMATIC_LINKS, ) async def non_href_links_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests with a page containing non-href links.""" await send_html_response( send, NON_HREF_LINKS, ) async def redirect_to_url(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests that should redirect to a specified full URL.""" query_params = get_query_params(scope.get('query_string', b'')) target_url = query_params.get('url', 'http://example.com') status_code = int(query_params.get('status', 302)) await send( { 'type': 'http.response.start', 'status': status_code, 'headers': [ [b'content-type', b'text/plain; charset=utf-8'], [b'location', target_url.encode()], ], } ) await send({'type': 'http.response.body', 'body': f'Redirecting to {target_url}...'.encode()}) async def echo_user_agent(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Echo back the user agent header as a response.""" headers = get_headers_dict(scope) user_agent = headers.get('user-agent', 'Not provided') await send_json_response(send, {'user-agent': user_agent}) async def get_echo(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Echo back GET request details similar to httpbin.org/get.""" path = scope.get('path', '') query_string = scope.get('query_string', b'') args = get_query_params(query_string) headers = get_headers_dict(scope) origin = scope.get('client', ('unknown', 0))[0] host = headers.get('host', 'localhost') scheme = headers.get('x-forwarded-proto', 'http') url = f'{scheme}://{host}{path}' if query_string: url += f'?{query_string}' response = { 'args': args, 'headers': headers, 'origin': origin, 'url': url, } await send_json_response(send, response) async def set_complex_cookies(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests to set specific cookies with various attributes.""" headers = [ [b'content-type', b'text/plain; charset=utf-8'], [b'set-cookie', b'basic=1; Path=/; HttpOnly; SameSite=Lax'], [b'set-cookie', b'withpath=2; Path=/html; SameSite=None'], [b'set-cookie', b'strict=3; Path=/; SameSite=Strict'], [b'set-cookie', b'secure=4; Path=/; HttpOnly; Secure; SameSite=Strict; Partitioned'], [b'set-cookie', b'short=5; Path=/;'], [b'set-cookie', b'domain=6; Path=/; Domain=.127.0.0.1;'], ] await send( { 'type': 'http.response.start', 'status': 200, 'headers': headers, } ) await send({'type': 'http.response.body', 'body': b'Cookies have been set!'}) async def echo_content(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Echo back content (plain text or base64) with specified content-type.""" query_params = get_query_params(scope.get('query_string', b'')) content = query_params.get('content', '') base64_content = query_params.get('base64', '') c_type = query_params.get('c_type', 'text/html; charset=utf-8') out_content = base64.b64decode(base64_content) if base64_content else content.encode() await send( { 'type': 'http.response.start', 'status': 200, 'headers': [[b'content-type', c_type.encode()]], } ) await send({'type': 'http.response.body', 'body': out_content}) async def robots_txt(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the robots.txt file.""" await send_html_response(send, ROBOTS_TXT) async def get_compressed(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Return large gzip compressed content.""" await send( { 'type': 'http.response.start', 'status': 200, 'headers': [[b'content-encoding', b'gzip']], } ) await send({'type': 'http.response.body', 'body': gzip.compress(HELLO_WORLD * 1000)}) async def slow_response(scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests with a configurable delay to test timeouts.""" query_params = get_query_params(scope.get('query_string', b'')) delay = float(query_params.get('delay', '5')) # Default 5 second delay await asyncio.sleep(delay) await send_html_response(send, HELLO_WORLD) async def infinite_scroll_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the infinite scroll page.""" await send_html_response( send, INFINITE_SCROLL, ) async def resource_loading_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the resource loading page.""" await send_html_response( send, RESOURCE_LOADING_PAGE, ) async def base_index_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the base index page.""" host = f'http://{get_headers_dict(_scope).get("host", "localhost")}' content = BASE_INDEX.format(host=host).encode() await send_html_response( send, content, ) async def start_enqueue_non_href_endpoint(_scope: dict[str, Any], _receive: Receive, send: Send) -> None: """Handle requests for the base index page.""" host = f'http://{get_headers_dict(_scope).get("host", "localhost")}' content = START_ENQUEUE_NON_HREF.format(host=host).encode() await send_html_response( send, content, ) class TestServer(Server): """A test HTTP server implementation based on Uvicorn Server.""" @property def url(self) -> URL: """Get the base URL of the server. Returns: A URL instance with the server's base URL. """ protocol = 'https' if self.config.is_ssl else 'http' return URL(f'{protocol}://{self.config.host}:{self.config.port}/') async def serve(self, sockets: list[socket] | None = None) -> None: """Run the server and set up restart capability. Args: sockets: Optional list of sockets to bind to. """ self.restart_requested = asyncio.Event() loop = asyncio.get_event_loop() tasks = { loop.create_task(super().serve(sockets=sockets)), loop.create_task(self.watch_restarts()), } await asyncio.wait(tasks) async def restart(self) -> None: """Request server restart and wait for it to complete. This method can be called from a different thread than the one the server is running on, and from a different async environment. """ self.started = False self.restart_requested.set() while not self.started: # noqa: ASYNC110 await asyncio.sleep(0.2) async def watch_restarts(self) -> None: """Watch for and handle restart requests.""" while True: if self.should_exit: return try: await asyncio.wait_for(self.restart_requested.wait(), timeout=0.1) except asyncio.TimeoutError: continue self.restart_requested.clear() await self.shutdown() await self.startup() def run(self, sockets: list[socket] | None = None) -> None: """Run the server.""" # Set the event loop policy in thread with server for Windows and Python 3.12+. # This is necessary because there are problems with closing connections when using `ProactorEventLoop`. if sys.version_info >= (3, 12) and sys.platform == 'win32': return asyncio.run(self.serve(sockets=sockets), loop_factory=asyncio.SelectorEventLoop) super().run(sockets=sockets) return None def serve_in_thread(server: TestServer) -> Iterator[TestServer]: """Run a server in a background thread and yield it.""" thread = threading.Thread(target=server.run, daemon=True) thread.start() try: # Bound the startup wait: when uvicorn fails to bind (e.g. port collision under xdist), the worker thread # exits without ever setting server.started, so an unbounded loop would hang until pytest-timeout kills # the suite 30 minutes later. deadline = time.monotonic() + 30 while not server.started: if not thread.is_alive(): raise RuntimeError('Test server thread exited before becoming ready (likely a bind failure).') if time.monotonic() > deadline: raise RuntimeError('Test server did not become ready within 30s.') time.sleep(1e-3) yield server finally: server.should_exit = True thread.join(timeout=10) if thread.is_alive(): # Uvicorn occasionally ignores should_exit; force_exit aborts the asyncio loop so teardown cannot hang # the suite indefinitely. server.force_exit = True thread.join(timeout=5)