crawlee-python/tests/unit/server.py

566 lines
19 KiB
Python

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"""<?xml version="1.0"?>
<hello>world</hello>""",
)
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)