diff --git a/README.md b/README.md index bf13fc9..0d2fc0e 100644 --- a/README.md +++ b/README.md @@ -92,6 +92,22 @@ res = session.post( ) ``` +Example 3 - Async: +```python +import asyncio +import tls_client + +async def main(): + # Takes the same arguments as tls_client.Session. + async with tls_client.AsyncSession(client_identifier="chrome_120") as session: + responses = await asyncio.gather( + *(session.get(f"https://www.example.com/{i}") for i in range(10)) + ) + print([res.status_code for res in responses]) + +asyncio.run(main()) +``` + # Pyinstaller / Pyarmor **If you want to pack the library with Pyinstaller or Pyarmor, make sure to add this to your command:** diff --git a/tests/integration/test_async_concurrency.py b/tests/integration/test_async_concurrency.py new file mode 100644 index 0000000..743f5ad --- /dev/null +++ b/tests/integration/test_async_concurrency.py @@ -0,0 +1,171 @@ +import asyncio +import json +import threading +import time +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +from tls_client import AsyncSession + +pytestmark = pytest.mark.integration + +SLOW_SECONDS = 0.3 + + +class Server(ThreadingHTTPServer): + # The default backlog of 5 makes the OS drop connections when 20 requests arrive at once. + request_queue_size = 64 + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), Handler) + self.lock = threading.Lock() + self.in_flight = 0 + self.peak_in_flight = 0 + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + +class Handler(BaseHTTPRequestHandler): + """/echo returns the request as JSON, /slow answers after SLOW_SECONDS, /cookie/ sets that cookie.""" + + server: Server + + def do_GET(self) -> None: + self.respond() + + def do_POST(self) -> None: + self.respond() + + def respond(self) -> None: + with self.server.lock: + self.server.in_flight += 1 + self.server.peak_in_flight = max(self.server.peak_in_flight, self.server.in_flight) + try: + body = self.rfile.read(int(self.headers.get("Content-Length") or 0)).decode() + self.send_response(200) + if self.path == "/slow": + time.sleep(SLOW_SECONDS) + elif self.path.startswith("/cookie/"): + name = self.path.removeprefix("/cookie/") + self.send_header("Set-Cookie", f"{name}={name}-value; Path=/") + payload = json.dumps( + {"method": self.command, "path": self.path, "headers": dict(self.headers), "body": body} + ).encode() + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + finally: + with self.server.lock: + self.server.in_flight -= 1 + + def log_message(self, format: str, *args: object) -> None: + pass + + +@pytest.fixture +def server() -> Iterator[Server]: + srv = Server() + thread = threading.Thread(target=srv.serve_forever, kwargs={"poll_interval": 0.05}, daemon=True) + thread.start() + yield srv + srv.shutdown() + srv.server_close() + thread.join() + + +def test_get_and_post_reach_the_server(server): + async def main(): + async with AsyncSession() as session: + got = await session.get(f"{server.url}/echo", params={"q": "1"}) + posted = await session.post(f"{server.url}/echo", json={"name": "tls"}) + return got, posted + + got, posted = asyncio.run(main()) + + assert got.status_code == 200 + assert got.json()["method"] == "GET" + assert got.json()["path"] == "/echo?q=1" + assert posted.json()["method"] == "POST" + assert json.loads(posted.json()["body"]) == {"name": "tls"} + assert posted.json()["headers"]["Content-Type"] == "application/json" + + +def test_constructor_arguments_reach_the_session(server): + async def main(): + async with AsyncSession(client_identifier="firefox_120", force_http1=True) as session: + return session, await session.get(f"{server.url}/echo") + + session, response = asyncio.run(main()) + + assert session.session.client_identifier == "firefox_120" + assert session.session.force_http1 is True + assert response.status_code == 200 + + +def test_session_headers_are_sent(server): + async def main(): + async with AsyncSession() as session: + session.headers["X-Session"] = "1" + return await session.get(f"{server.url}/echo", headers={"X-Request": "2"}) + + received = asyncio.run(main()).json()["headers"] + + assert received["X-Session"] == "1" + assert received["X-Request"] == "2" + + +def test_concurrent_requests_overlap(server): + count = 20 + + async def main(): + with ThreadPoolExecutor(max_workers=count) as executor: + async with AsyncSession(executor=executor) as session: + start = time.perf_counter() + responses = await asyncio.gather(*(session.get(f"{server.url}/slow") for _ in range(count))) + return responses, time.perf_counter() - start + + responses, elapsed = asyncio.run(main()) + + assert [r.status_code for r in responses] == [200] * count + # Serially this takes count * SLOW_SECONDS = 6s. + assert elapsed < count * SLOW_SECONDS / 3 + assert server.peak_in_flight > count / 2 + + +def test_cookies_from_concurrent_responses_all_end_up_in_the_jar(server): + names = [f"c{i}" for i in range(20)] + + async def main(): + with ThreadPoolExecutor(max_workers=len(names)) as executor: + async with AsyncSession(executor=executor) as session: + await asyncio.gather(*(session.get(f"{server.url}/cookie/{name}") for name in names)) + sent_back = await session.get(f"{server.url}/echo") + return session.cookies.get_dict(), sent_back + + jar, sent_back = asyncio.run(main()) + + assert jar == {name: f"{name}-value" for name in names} + sent_cookies = dict(pair.split("=") for pair in sent_back.json()["headers"]["Cookie"].split("; ")) + assert sent_cookies == jar + + +def test_concurrent_bodies_keep_their_own_content_type(server): + async def main(): + async with AsyncSession() as session: + responses = await asyncio.gather( + *(session.post(f"{server.url}/echo", json={"i": i}) for i in range(5)), + *(session.post(f"{server.url}/echo", data={"i": i}) for i in range(5)), + ) + return session, responses + + session, responses = asyncio.run(main()) + + content_types = [r.json()["headers"]["Content-Type"] for r in responses] + assert content_types == ["application/json"] * 5 + ["application/x-www-form-urlencoded"] * 5 + assert "Content-Type" not in session.headers diff --git a/tests/unit/test_async_sessions.py b/tests/unit/test_async_sessions.py new file mode 100644 index 0000000..195b329 --- /dev/null +++ b/tests/unit/test_async_sessions.py @@ -0,0 +1,223 @@ +import asyncio +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any + +import pytest + +from tls_client import AsyncSession +from tls_client.cookies import cookiejar_from_dict +from tls_client.exceptions import TLSClientException +from tls_client.structures import CaseInsensitiveDict + +URL = "https://example.com/path" + + +def sent_cookies(payload: dict) -> dict[str, str]: + return {cookie["name"]: cookie["value"] for cookie in payload["requestCookies"]} + + +class CountingExecutor(ThreadPoolExecutor): + def __init__(self) -> None: + super().__init__(max_workers=1) + self.submitted = 0 + + def submit(self, fn: Any, /, *args: Any, **kwargs: Any) -> Future: + self.submitted += 1 + return super().submit(fn, *args, **kwargs) + + +class TestSettings: + def test_constructor_arguments_are_forwarded_to_the_session(self, fake_lib): + async def main(): + async with AsyncSession("firefox_120", random_tls_extension_order=True) as session: + await session.get(URL) + + asyncio.run(main()) + + assert fake_lib.sent["tlsClientIdentifier"] == "firefox_120" + assert fake_lib.sent["withRandomTLSExtensionOrder"] is True + + def test_properties_read_and_write_the_wrapped_session(self, fake_lib): + session = AsyncSession() + jar = cookiejar_from_dict({"a": "1"}) + + session.headers = CaseInsensitiveDict({"X-Test": "1"}) + session.cookies = jar + session.proxies = "http://proxy:8080" + session.params = {"q": "1"} + session.timeout_seconds = 5 + + assert session.session.headers == {"X-Test": "1"} + assert session.session.cookies is jar + assert session.session.proxies == "http://proxy:8080" + assert session.session.params == {"q": "1"} + assert session.session.timeout_seconds == 5 + assert session.headers is session.session.headers + assert session.cookies is jar + assert session.proxies == "http://proxy:8080" + assert session.params == {"q": "1"} + assert session.timeout_seconds == 5 + + def test_session_settings_are_sent(self, fake_lib): + async def main(): + async with AsyncSession() as session: + session.headers["X-Session"] = "1" + session.cookies.set("from_session", "1") + session.proxies = "http://proxy:8080" + session.timeout_seconds = 5 + session.session.force_http1 = True + await session.get(URL) + + asyncio.run(main()) + + assert fake_lib.sent["headers"]["X-Session"] == "1" + assert sent_cookies(fake_lib.sent) == {"from_session": "1"} + assert fake_lib.sent["proxyUrl"] == "http://proxy:8080" + assert fake_lib.sent["timeoutMilliseconds"] == 5000 + assert fake_lib.sent["forceHttp1"] is True + + +class TestRequests: + @pytest.mark.parametrize("method", ["GET", "OPTIONS", "HEAD", "POST", "PUT", "PATCH", "DELETE"]) + def test_verb_helper_sends_its_method(self, fake_lib, method): + async def main(): + async with AsyncSession() as session: + return await getattr(session, method.lower())(URL) + + response = asyncio.run(main()) + + assert fake_lib.sent["requestMethod"] == method + assert response.status_code == 200 + + @pytest.mark.parametrize("method", ["post", "put", "patch"]) + def test_body_verbs_take_data_as_second_positional_argument(self, fake_lib, method): + async def main(): + async with AsyncSession() as session: + await getattr(session, method)(URL, "raw") + + asyncio.run(main()) + + assert fake_lib.sent["requestBody"] == "raw" + + def test_keyword_arguments_are_forwarded(self, fake_lib): + async def main(): + async with AsyncSession() as session: + await session.execute_request("GET", URL, params={"q": "1"}, allow_redirects=True) + + asyncio.run(main()) + + assert fake_lib.sent["requestUrl"] == f"{URL}?q=1" + assert fake_lib.sent["followRedirects"] is True + + def test_requests_use_the_native_session_of_the_wrapped_session(self, fake_lib): + async def main(): + async with AsyncSession() as session: + await asyncio.gather(session.get(URL), session.get(URL)) + + asyncio.run(main()) + + session_ids = {payload["sessionId"] for payload in fake_lib.payloads} + assert session_ids == set(fake_lib.destroyed_session_ids) + assert len(session_ids) == 1 + + def test_library_error_is_raised_as_tls_client_exception(self, fake_lib): + fake_lib.reply_with(status=0, body="connection refused", headers={"Set-Cookie": ["a=1; Path=/"]}) + + async def main(): + async with AsyncSession() as session: + with pytest.raises(TLSClientException, match="connection refused"): + await session.get(URL) + return session + + assert len(asyncio.run(main()).cookies) == 0 + + def test_requests_run_on_the_given_executor(self, fake_lib): + with CountingExecutor() as executor: + + async def main(): + async with AsyncSession(executor=executor) as session: + await session.get(URL) + await session.post(URL) + + asyncio.run(main()) + + # Two requests, then close. + assert executor.submitted == 3 + + +class TestState: + def test_cookies_from_the_response_are_kept_for_the_next_request(self, fake_lib): + fake_lib.reply_with(headers={"Set-Cookie": ["a=1; Path=/", "b=2; Path=/"]}) + + async def main(): + async with AsyncSession() as session: + response = await session.get(URL) + await session.get(URL) + return session, response + + session, response = asyncio.run(main()) + + assert dict(response.cookies) == {"a": "1", "b": "2"} + assert dict(session.cookies) == {"a": "1", "b": "2"} + assert sent_cookies(fake_lib.sent) == {"a": "1", "b": "2"} + + @pytest.mark.parametrize("cookies", [{"from_request": "2"}, cookiejar_from_dict({"from_request": "2"})]) + def test_request_cookies_are_kept_in_the_jar_like_session_does(self, fake_lib, cookies): + async def main(): + async with AsyncSession() as session: + session.cookies.set("from_session", "1") + await session.get(URL, cookies=cookies) + return session + + session = asyncio.run(main()) + + assert sent_cookies(fake_lib.sent) == {"from_session": "1", "from_request": "2"} + assert dict(session.cookies) == {"from_session": "1", "from_request": "2"} + + def test_inferred_content_type_applies_to_one_request_only(self, fake_lib): + async def main(): + async with AsyncSession() as session: + await session.post(URL, json={"a": 1}) + await session.post(URL, data={"a": "1"}) + return session + + session = asyncio.run(main()) + + assert fake_lib.payloads[0]["headers"]["Content-Type"] == "application/json" + assert fake_lib.payloads[1]["headers"]["Content-Type"] == "application/x-www-form-urlencoded" + assert "Content-Type" not in session.headers + + def test_request_headers_leave_session_headers_untouched(self, fake_lib): + async def main(): + async with AsyncSession() as session: + before = session.headers.copy() + await session.get(URL, headers={"X-Request": "1", "Accept": None}) + return before, session.headers + + before, after = asyncio.run(main()) + + assert after == before + + +class TestLifecycle: + def test_close_destroys_the_native_session(self, fake_lib): + async def main(): + session = AsyncSession() + await session.get(URL) + return await session.close() + + result = asyncio.run(main()) + + assert fake_lib.destroyed_session_ids == [fake_lib.sent["sessionId"]] + assert f"destroy-{fake_lib.sent['sessionId']}" in fake_lib.freed_ids + assert '"success": true' in result + + def test_leaving_the_async_with_block_closes_the_session_even_on_error(self, fake_lib): + async def main(): + async with AsyncSession(): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + asyncio.run(main()) + + assert len(fake_lib.destroyed_session_ids) == 1 diff --git a/tls_client/__init__.py b/tls_client/__init__.py index ef5d1be..fe37686 100644 --- a/tls_client/__init__.py +++ b/tls_client/__init__.py @@ -12,6 +12,7 @@ # tls-client: https://github.com/bogdanfinn/tls-client # requests: https://github.com/psf/requests +from .async_sessions import AsyncSession from .sessions import Session -__all__ = ["Session"] +__all__ = ["AsyncSession", "Session"] diff --git a/tls_client/async_sessions.py b/tls_client/async_sessions.py new file mode 100644 index 0000000..a8d5a11 --- /dev/null +++ b/tls_client/async_sessions.py @@ -0,0 +1,139 @@ +import asyncio +import copy +import functools +from collections.abc import Callable +from concurrent.futures import Executor +from typing import Any, TypeVar + +from .cookies import RequestsCookieJar +from .response import Response +from .sessions import Session, _Data, _Json, _Proxy +from .structures import CaseInsensitiveDict + +_T = TypeVar("_T") + + +class AsyncSession: + """asyncio front end for :class:`Session`. + + Takes the same arguments as ``Session`` and runs each request in a worker thread, so requests on one session can + run concurrently. Use it from a single event loop. + """ + + def __init__(self, *args: Any, executor: Executor | None = None, **kwargs: Any) -> None: + self._session = Session(*args, **kwargs) + # Runs the blocking calls; None means the event loop's default thread pool, whose size limits how many requests + # run at once. + self._executor = executor + + # --- Session state ------------------------------------------------------------------------------------------------ + # Changes apply to requests started afterwards. Requests that are already running keep what they started with. + + @property + def session(self) -> Session: + """The wrapped Session, for the settings that have no property here.""" + return self._session + + @property + def headers(self) -> CaseInsensitiveDict[str]: + return self._session.headers + + @headers.setter + def headers(self, value: CaseInsensitiveDict[str]) -> None: + self._session.headers = value + + @property + def cookies(self) -> RequestsCookieJar: + return self._session.cookies + + @cookies.setter + def cookies(self, value: RequestsCookieJar) -> None: + self._session.cookies = value + + @property + def proxies(self) -> _Proxy: + return self._session.proxies + + @proxies.setter + def proxies(self, value: _Proxy) -> None: + self._session.proxies = value + + @property + def params(self) -> dict[str, Any]: + return self._session.params + + @params.setter + def params(self, value: dict[str, Any]) -> None: + self._session.params = value + + @property + def timeout_seconds(self) -> float: + return self._session.timeout_seconds + + @timeout_seconds.setter + def timeout_seconds(self, value: float) -> None: + self._session.timeout_seconds = value + + # --- Lifecycle ---------------------------------------------------------------------------------------------------- + + async def __aenter__(self) -> "AsyncSession": + return self + + async def __aexit__(self, *args: object) -> None: + await self.close() + + async def close(self) -> str: + """Frees the native session. Await any running requests first.""" + return await self._run(self._session.close) + + # --- Requests ----------------------------------------------------------------------------------------------------- + + async def execute_request(self, method: str, url: str, **kwargs: Any) -> Response: + # Session.execute_request reads the session's headers and writes to its cookie jar, and neither is safe to use + # from several threads at once. Each request therefore runs on a shallow copy of the session: it shares the + # native session, so connections are reused, but gets its own headers and cookie jar. + request_session = copy.copy(self._session) + request_session.headers = CaseInsensitiveDict(self._session.headers) + request_session.cookies = self._session.cookies.copy() + + response = await self._run(request_session.execute_request, method=method, url=url, **kwargs) + + # Back on the event loop thread, so no other request touches the jar while the new cookies go in. Like Session, + # keep the cookies sent with the request, then the ones the server set. + cookies = kwargs.get("cookies") + if cookies: + self._session.cookies.update(cookies) + self._session.cookies.update(response.cookies) + return response + + async def get(self, url: str, **kwargs: Any) -> Response: + """Sends a GET request""" + return await self.execute_request(method="GET", url=url, **kwargs) + + async def options(self, url: str, **kwargs: Any) -> Response: + """Sends a OPTIONS request""" + return await self.execute_request(method="OPTIONS", url=url, **kwargs) + + async def head(self, url: str, **kwargs: Any) -> Response: + """Sends a HEAD request""" + return await self.execute_request(method="HEAD", url=url, **kwargs) + + async def post(self, url: str, data: _Data | None = None, json: _Json | None = None, **kwargs: Any) -> Response: + """Sends a POST request""" + return await self.execute_request(method="POST", url=url, data=data, json=json, **kwargs) + + async def put(self, url: str, data: _Data | None = None, json: _Json | None = None, **kwargs: Any) -> Response: + """Sends a PUT request""" + return await self.execute_request(method="PUT", url=url, data=data, json=json, **kwargs) + + async def patch(self, url: str, data: _Data | None = None, json: _Json | None = None, **kwargs: Any) -> Response: + """Sends a PATCH request""" + return await self.execute_request(method="PATCH", url=url, data=data, json=json, **kwargs) + + async def delete(self, url: str, **kwargs: Any) -> Response: + """Sends a DELETE request""" + return await self.execute_request(method="DELETE", url=url, **kwargs) + + async def _run(self, func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(self._executor, functools.partial(func, *args, **kwargs))