diff --git a/README.md b/README.md index 1793e2fd..c5792520 100644 --- a/README.md +++ b/README.md @@ -101,26 +101,26 @@ See [the devtools reference][devtools-ref] for a list of possible commands. Try adding the following to the example shown above: ```python - # Callback for printing result - async def dump_event(response): - print(str(response)) - - - # Callback for raising result as error - async def error_event(response): - raise Exception(str(response)) - - - browser.subscribe("Target.targetCrashed", error_event) - new_tab.subscribe("Page.loadEventFired", dump_event) - browser.subscribe("Target.*", dump_event) # dumps all "Target" events - response = await new_tab.subscribe_once("Page.lifecycleEvent") - # do something with response - browser.unsubscribe("Target.*") - # events are always sent to a browser or tab, - # but the documentation isn't always clear which. - # Dumping all: `browser.subscribe("*", dump_event)` (on tab too) - # can be useful (but verbose) for debugging. +# Callback for printing result +async def dump_event(response): + print(str(response)) + + +# Callback for raising result as error +async def error_event(response): + raise Exception(str(response)) + + +browser.subscribe("Target.targetCrashed", error_event) +new_tab.subscribe("Page.loadEventFired", dump_event) +browser.subscribe("Target.*", dump_event) # dumps all "Target" events +response = await new_tab.subscribe_once("Page.lifecycleEvent") +# do something with response +browser.unsubscribe("Target.*") +# events are always sent to a browser or tab, +# but the documentation isn't always clear which. +# Dumping all: `browser.subscribe("*", dump_event)` (on tab too) +# can be useful (but verbose) for debugging. ``` ## Synchronous Use @@ -128,7 +128,7 @@ Try adding the following to the example shown above: You can use this library without `asyncio`, ```python -my_browser = choreo.Browser() # blocking until open +my_browser = choreo.Browser() # blocking until open ``` However, you must call `browser.pipe.read_jsons(blocking=True|False)` manually, diff --git a/pyproject.toml b/pyproject.toml index 02ca206e..f4d18cd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -81,12 +81,14 @@ src = ["src"] select = ["ALL"] ignore = [ "ANN", # no types - "EM", # allow strings in raise(), despite python being ugly about it - "TRY003", # allow long error messages inside raise() + "COM812", # manual says linter rule conflicts with formatter + "CPY001", # Don't require a copyright notice at the top of a file "D203", # No blank before class docstring (D211 = require blank line) "D212", # Commit message style docstring is D213, ignore D212 - "COM812", # manual says linter rule conflicts with formatter + "EM", # allow strings in raise(), despite python being ugly about it + "G004", # fstrings in my logs "ISC001", # manual says litner rule conflicts with formatter + "PT003", # scope="function" implied but I like readability "RET504", # Allow else if unnecessary because more readable "RET505", # Allow else if unnecessary because more readable "RET506", # Allow else if unnecessary because more readable @@ -94,8 +96,7 @@ ignore = [ "RET508", # Allow else if unnecessary because more readable "RUF012", # We don't do typing, so no typing "SIM105", # Too opionated (try-except-pass) - "PT003", # scope="function" implied but I like readability - "G004", # fstrings in my logs + "TRY003", # allow long error messages inside raise() ] [tool.ruff.lint.per-file-ignores] @@ -112,6 +113,9 @@ asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" log_cli = false addopts = "--import-mode=append" +markers = [ + "slow: moves enough data to be worth skipping by default (`-m 'not slow'`)", +] # tell poe to use the env we give it, otherwise it detects uv and overrides flags [tool.poe] @@ -119,7 +123,10 @@ executor.type = "simple" [tool.poe.tasks] test_proc = "pytest --log-level=1 -W error -n auto -v -rfE --capture=fd tests/test_process.py" -test_fn = "pytest --log-level=1 -W error -n auto -v -rfE --capture=fd --ignore=tests/test_process.py" +test_fn = "pytest --log-level=1 -W error -n auto -v -rfE --capture=fd -m 'not slow' --ignore=tests/test_process.py" +# Skip `--log-level=1` and `-n auto` to allow slow tests to complete without +# formatting wall of text and to avoid tests using up all memory (~1GB per thread) +test_slow = "pytest -W error -v -rfE --capture=fd -m slow --ignore=tests/test_process.py" debug-test_proc = "pytest --log-level=1 -W error -vvvx -rA --show-capture=no --capture=no tests/test_process.py" debug-test_fn = "pytest --log-level=1 -W error -vvvx -rA --show-capture=no --capture=no --ignore=tests/test_process.py" diff --git a/src/choreographer/channels/__init__.py b/src/choreographer/channels/__init__.py index 1257759a..7e4ccbe2 100644 --- a/src/choreographer/channels/__init__.py +++ b/src/choreographer/channels/__init__.py @@ -5,7 +5,12 @@ """ -from ._errors import BlockWarning, ChannelClosedError, JSONError +from ._errors import ( + BlockWarning, + ChannelClosedError, + JSONError, + MessageTooLargeError, +) from ._wire import register_custom_encoder from .pipe import Pipe @@ -13,6 +18,7 @@ "BlockWarning", "ChannelClosedError", "JSONError", + "MessageTooLargeError", "Pipe", "register_custom_encoder", ] diff --git a/src/choreographer/channels/_errors.py b/src/choreographer/channels/_errors.py index 48baf613..1114dc87 100644 --- a/src/choreographer/channels/_errors.py +++ b/src/choreographer/channels/_errors.py @@ -1,3 +1,6 @@ +from __future__ import annotations + + class BlockWarning(UserWarning): """A warning for when block modification operations used on incompatible OS.""" @@ -8,3 +11,43 @@ class ChannelClosedError(IOError): class JSONError(RuntimeError): """Another JSONError.""" + + +class MessageTooLargeError(RuntimeError): + """ + An error for when a message won't fit in the browser's receive buffer. + + The browser closes the connection outright if we write a message bigger + than its buffer, so we refuse to write it and raise this instead. + """ + + size: int + """The size, in bytes, of the message we refused to write.""" + max_size: int + """The largest message the browser will accept.""" + payload: str | None + """ + The serialized message. + + It is kept so that callers who know how to break the message up can + reuse it instead of serializing all over again. It is deliberately left + out of the error text: it can be hundreds of megabytes. + """ + + def __init__(self, size: int, max_size: int, payload: str | None = None) -> None: + """ + Construct a MessageTooLargeError. + + Args: + size: the size of the message in bytes. + max_size: the largest message the browser will accept. + payload: the serialized message, if the caller should have it. + + """ + super().__init__( + f"Message is {size} bytes, which is over the browser's " + f"{max_size} byte limit. It was not sent.", + ) + self.size = size + self.max_size = max_size + self.payload = payload diff --git a/src/choreographer/channels/_wire.py b/src/choreographer/channels/_wire.py index 73b3acf5..6c2c7197 100644 --- a/src/choreographer/channels/_wire.py +++ b/src/choreographer/channels/_wire.py @@ -41,7 +41,17 @@ def default(self, o: Any) -> Any: return simplejson.JSONEncoder.default(self, o) -def serialize(obj: Any) -> bytes: +def serialize_str(obj: Any) -> str: + """ + Serialize an object to a JSON string. + + Use `serialize()` to return a value that's ready to send to Chrome. + This exists for callers that need to slice the string up first. + + Args: + obj: Any Python object that serializes to JSON. + + """ try: if not _custom_encoder: message = simplejson.dumps( @@ -57,7 +67,11 @@ def serialize(obj: Any) -> bytes: _logger.debug(f"Serialized: {message[:15]}...{message[-15:]}, size: {len(message)}") _logger.debug2(f"Whole message: {message}") - return message.encode("utf-8") + return message + + +def serialize(obj: Any) -> bytes: + return serialize_str(obj).encode("utf-8") def deserialize(message: str) -> Any: diff --git a/src/choreographer/channels/pipe.py b/src/choreographer/channels/pipe.py index f7f22cf8..5a03bfac 100644 --- a/src/choreographer/channels/pipe.py +++ b/src/choreographer/channels/pipe.py @@ -13,7 +13,12 @@ import logistro from . import _wire as wire -from ._errors import BlockWarning, ChannelClosedError, JSONError +from ._errors import ( + BlockWarning, + ChannelClosedError, + JSONError, + MessageTooLargeError, +) if TYPE_CHECKING: from typing import Any, Mapping, Sequence @@ -24,6 +29,14 @@ _logger = logistro.getLogger(__name__) +MAX_MESSAGE_SIZE = 100 * 1024 * 1024 +""" +The biggest message Chrome will read off the pipe, in bytes. + +This mirrors `kReceiveBufferSizeForDevTools` in Chrome's +`content/browser/devtools/devtools_pipe_handler.cc`. +""" + # should be closing my ends from the start? @@ -84,7 +97,20 @@ def write_json(self, obj: Mapping[str, Any]) -> tuple[float, float]: Send one json down the pipe. Args: - obj: any python object that serializes to json. + obj: Any python object that serializes to JSON. + + Raises: + ChannelClosedError: If the pipe was never opened or is already + closed, or if the OS write fails. A failed write closes the + pipe, so nothing can be sent after this. + MessageTooLargeError: If the message won't fit in Chrome's buffer. + Nothing is written, so the channel is still good afterwards. + The error carries the serialized message so that callers who + can break it up don't have to serialize it a second time. + TypeError: If `obj` contains something the encoder doesn't know how + to turn into JSON. + UnicodeEncodeError: If the serialized message contains lone + surrogates, which have no UTF-8 representation. """ if not self.is_ready(): @@ -92,7 +118,15 @@ def write_json(self, obj: Mapping[str, Any]) -> tuple[float, float]: "The communication channel was either never " "opened or closed. Was .open() or .close() called?", ) - encoded_message = wire.serialize(obj) + b"\0" + message = wire.serialize_str(obj) + encoded_message = message.encode("utf-8") + b"\0" + if len(encoded_message) > MAX_MESSAGE_SIZE: + # Don't close(): we haven't written anything, the pipe is fine. + raise MessageTooLargeError( + len(encoded_message), + MAX_MESSAGE_SIZE, + payload=message, + ) _logger.debug( f"Writing message {encoded_message[:15]!r}...{encoded_message[-15:]!r}, " f"size: {len(encoded_message)}.", diff --git a/src/choreographer/cli/_cli_utils.py b/src/choreographer/cli/_cli_utils.py index 5b02fcc5..6d87be37 100644 --- a/src/choreographer/cli/_cli_utils.py +++ b/src/choreographer/cli/_cli_utils.py @@ -131,7 +131,7 @@ def get_chrome_sync( # noqa: C901, PLR0912, PLR0915 if i: _logger.info("Loading chrome from list") - raw_json = urllib.request.urlopen( # noqa: S310 audit url for schemes + raw_json = urllib.request.urlopen( _chrome_for_testing_url, ).read() browser_list = json.loads( diff --git a/src/choreographer/protocol/_chunking.py b/src/choreographer/protocol/_chunking.py new file mode 100644 index 00000000..9a270b5e --- /dev/null +++ b/src/choreographer/protocol/_chunking.py @@ -0,0 +1,251 @@ +""" +Break up `Runtime.callFunctionOn` commands that are too big for the pipe. + +Chrome reads one complete JSON message at a time off the devtools pipe, and +it won't read one bigger than 100MB (see `channels.pipe.MAX_MESSAGE_SIZE`). +The devtools protocol has no way to split a message, so in general an +oversized command is simply an error. + +`Runtime.callFunctionOn` is the exception. We can stash pieces of the +message in a JavaScript array on the page, then run a function that glues +them back together. That's what this module does. + +None of this is part of the public API. `Session.send_command` falls back to +it on its own, and the function you asked Chrome to run can't tell the +difference: it still receives the same parsed arguments it always would. +""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import logistro + +from . import DevtoolsProtocolError + +if TYPE_CHECKING: + from typing import Any, MutableMapping + + from . import BrowserCommand, BrowserResponse + from .devtools_async import Session + +_logger = logistro.getLogger(__name__) + +CHUNK_SIZE = 10 * 1024 * 1024 +""" +How much of the message to put in each piece, in characters. + +Deliberately well under `MAX_MESSAGE_SIZE` to avoid issues with +high byte count UTF-8 characters and escaped characters. +""" + +_CHUNKABLE_METHOD = "Runtime.callFunctionOn" +"""The one command we know how to break up.""" + +_HELPER_METHOD = "Runtime.callFunctionOn" +""" +How we send our own setup, push, and cleanup calls. + +This is always `Runtime.callFunctionOn`, whatever the oversized command we +were handed happens to be. It only matches `_CHUNKABLE_METHOD` today because +that is the one method we can break up. +""" + +_STORE = "window.__choreo_chunks" + +# `arguments` entries can also be `objectId` or `unserializableValue` +# handles, which refer to things that only exist in the browser. Those can't +# be rebuilt from the text of the message, so we only take the plain ones. +_VALUE_KEY = "value" + +_INIT_FN = f"function(k){{ ({_STORE} = {_STORE} || {{}})[k] = []; }}" + +_PUSH_FN = f"function(k, c){{ {_STORE}[k].push(c); }}" + +_DELETE_FN = f"function(k){{ if ({_STORE}) {{ delete {_STORE}[k]; }} }}" + +_counter = itertools.count() + + +def is_chunkable(command: BrowserCommand) -> bool: + """ + Report whether we know how to break this command up. + + Args: + command: The command that was too big to send + + """ + method = command.get("method") + if method != _CHUNKABLE_METHOD: + _logger.debug(f"Can't chunk {method}: only {_CHUNKABLE_METHOD} can be.") + return False + params = command.get("params") or {} + if params.get("throwOnSideEffect"): + _logger.debug( + "Can't chunk a command with throwOnSideEffect: breaking it up means " + "writing to the page, which is what the flag forbids.", + ) + return False + if not params.get("functionDeclaration"): + _logger.debug("Can't chunk: no functionDeclaration to wrap.") + return False + arguments = params.get("arguments") or [] + # If even one argument is a browser-side handle we can't rebuild the call + # from the message text, so we don't try + if not arguments or not all( + isinstance(argument, dict) and _VALUE_KEY in argument for argument in arguments + ): + _logger.debug( + "Can't chunk: every argument has to be a plain " + f"{_VALUE_KEY!r}, and at least one has to exist.", + ) + return False + return True + + +def _build_wrapper(user_fn: str) -> str: + """ + Wrap the caller's function in one that rebuilds its arguments first. + + What we glue back together is the whole original command, not just the + big argument, because that is exactly what we already had serialized. + The extra envelope is a few hundred bytes on top of a very large + message, and pulling the arguments out of it in JavaScript is cheap. + """ + return ( + "function(k){" + "try{" + f"var cmd = JSON.parse({_STORE}[k].join(''));" + "var args = cmd.params.arguments.map(function(a){ return a.value; });" + f"return ({user_fn}).apply(this, args);" + "}finally{" + f"delete {_STORE}[k];" + "}" + "}" + ) + + +def _raise_for_error(response: BrowserResponse) -> BrowserResponse: + if "error" in response: + raise DevtoolsProtocolError(response) + return response + + +def _js_error_text(details: MutableMapping[str, Any]) -> str: + """Pull the readable part out of a `Runtime.exceptionDetails`.""" + # `text` is usually just "Uncaught", so check exception + exception = details.get("exception") or {} + + return exception.get("description") or details.get("text") or "unknown error" + + +async def _send( + session: Session, + params: MutableMapping[str, Any], +) -> BrowserResponse: + """Send one piece, going around the too-big fallback so we can't recurse.""" + response = _raise_for_error( + await session._send_no_retry(_HELPER_METHOD, params), # noqa: SLF001 + ) + # Check for JS error response since that won't throw a top-level error + exception_details = response.get("result", {}).get("exceptionDetails") + if exception_details: + raise RuntimeError( + f"Chunked send failed in the page: {_js_error_text(exception_details)}", + ) + return response + + +async def send_chunked( + session: Session, + command: BrowserCommand, + payload: str, +) -> tuple[BrowserResponse, BrowserCommand]: + """ + Send an oversized `Runtime.callFunctionOn` in pieces. + + Args: + session: The session the original command was sent on. + command: The original command, too big to send in one go. + payload: The already serialized command, from the error the pipe + raised. We slice this rather than serializing all over again. + + Returns: + The response to the final call, and the command that produced it + (the caller needs it to look up timings). + + """ + params: MutableMapping[str, Any] = dict(command["params"]) + # Grab all potential identifiers for the execution context + execution_context_params = { + key: params[key] + for key in ("executionContextId", "objectId", "uniqueContextId") + if key in params + } + # Two calls in one page must not share a store, or they'd eat each other. + key = f"{session.session_id or 'browser'}:{next(_counter)}" + + # Use ceiling division to ensure whole number of chunks + n_chunks = -(-len(payload) // CHUNK_SIZE) + _logger.info( + f"Message too big for one write, sending it as {n_chunks} pieces " + f"under key {key}.", + ) + + try: + await _send( + session, + { + **execution_context_params, + "functionDeclaration": _INIT_FN, + "arguments": [{"value": key}], + }, + ) + for start in range(0, len(payload), CHUNK_SIZE): + await _send( + session, + { + **execution_context_params, + "functionDeclaration": _PUSH_FN, + "arguments": [ + {"value": key}, + {"value": payload[start : start + CHUNK_SIZE]}, + ], + }, + ) + + params["functionDeclaration"] = _build_wrapper( + command["params"]["functionDeclaration"], + ) + params["arguments"] = [{"value": key}] + # Everything else the caller asked for (awaitPromise, returnByValue, + # silent, userGesture, the execution context) rides along untouched. + final_command = session._build_command( # noqa: SLF001 + _HELPER_METHOD, + params, + ) + response = await session._send_built(final_command) # noqa: SLF001 + except Exception: + await _cleanup(session, execution_context_params, key) + raise + return response, final_command + + +async def _cleanup( + session: Session, + execution_context_params: MutableMapping[str, Any], + key: str, +) -> None: + """Drop the store if we bailed out before the wrapper could.""" + try: + await _send( + session, + { + **execution_context_params, + "functionDeclaration": _DELETE_FN, + "arguments": [{"value": key}], + }, + ) + except Exception: # noqa: BLE001 we're already failing, don't make it worse + _logger.debug(f"Couldn't clean up chunk store {key}.") diff --git a/src/choreographer/protocol/devtools_async.py b/src/choreographer/protocol/devtools_async.py index dd2d3a55..09516345 100644 --- a/src/choreographer/protocol/devtools_async.py +++ b/src/choreographer/protocol/devtools_async.py @@ -8,6 +8,9 @@ import logistro from choreographer import protocol +from choreographer.channels import MessageTooLargeError + +from . import _chunking if TYPE_CHECKING: import asyncio @@ -98,8 +101,43 @@ async def send_command( A message key (session, message id) tuple or None (Optional) A tuple[float, float, float] representing perf_counters() for write start, end, and read end. + On a chunked command, this covers only the final message, not + the pieces that carried the payload, so it reads as far quicker + than the call really was. + + Raises: + MessageTooLargeError: If the message is too big for Chrome's + buffer and isn't a `Runtime.callFunctionOn` we can break up. """ + json_command = self._build_command(command, params) + _logger.debug( + f"Cmd '{command}', param keys '{params.keys() if params else ''}', " + f"sessionId '{self.session_id}'", + ) + _logger.debug2(f"Full params: {str(params).replace('%', '%%')}") + try: + response = await self._send_built(json_command) + except MessageTooLargeError as e: + if e.payload is None or not _chunking.is_chunkable(json_command): + raise + # Break the command up into chunks so they can be sent piece by piece + # and reassmbled without hitting the CDP limit + response, json_command = await _chunking.send_chunked( + self, + json_command, + e.payload, + ) + if with_perf: + return (response, self._broker.get_perf(json_command)) + return response + + def _build_command( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + ) -> protocol.BrowserCommand: + """Give a command an id and wrap it up for the browser.""" current_id = self.message_id self.message_id += 1 json_command = protocol.BrowserCommand( @@ -113,18 +151,23 @@ async def send_command( json_command["sessionId"] = self.session_id if params: json_command["params"] = params - _logger.debug( - f"Cmd '{command}', param keys '{params.keys() if params else ''}', " - f"sessionId '{self.session_id}'", - ) - _logger.debug2(f"Full params: {str(params).replace('%', '%%')}") - if with_perf: - return ( - await self._broker.write_json(json_command), - self._broker.get_perf(json_command), - ) + return json_command + + async def _send_built( + self, + json_command: protocol.BrowserCommand, + ) -> protocol.BrowserResponse: + """Write a command with no too-big fallback, so we can't recurse.""" return await self._broker.write_json(json_command) + async def _send_no_retry( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + ) -> protocol.BrowserResponse: + """Build and write a command with no too-big fallback.""" + return await self._send_built(self._build_command(command, params)) + def subscribe( self, string: str, diff --git a/src/choreographer/protocol/devtools_async_helpers.py b/src/choreographer/protocol/devtools_async_helpers.py index dafbb3ba..c1d2d250 100644 --- a/src/choreographer/protocol/devtools_async_helpers.py +++ b/src/choreographer/protocol/devtools_async_helpers.py @@ -20,7 +20,7 @@ # racey. Optimistically, it's buffered and fired after subscription # even if the event happened in the past. # Doesn't seem to always work out that way, so we also use -# javascript to create a "loaded" event, but for the case +# JavaScript to create a "loaded" event, but for the case # where we need to timeout- loading a page that never resolves, # the browser might actually load an about:blank instead and then # fire the event, misleading the user, so we check the url. diff --git a/src/choreographer/protocol/devtools_sync.py b/src/choreographer/protocol/devtools_sync.py index fde81a1c..7860a1b4 100644 --- a/src/choreographer/protocol/devtools_sync.py +++ b/src/choreographer/protocol/devtools_sync.py @@ -64,6 +64,11 @@ def send_command( Returns: A message key (session, message id) tuple or None + Raises: + MessageTooLargeError: If the message is too big for Chrome's + buffer. Use the async `Session` if you need to send something + that big. + """ current_id = self.message_id self.message_id += 1 diff --git a/src/choreographer/utils/_tmpfile.py b/src/choreographer/utils/_tmpfile.py index f27928ad..0c7ca52d 100644 --- a/src/choreographer/utils/_tmpfile.py +++ b/src/choreographer/utils/_tmpfile.py @@ -183,7 +183,7 @@ def remove_readonly( if hasattr(self, "temp_dir"): del self.temp_dir _logger.info("shutil.rmtree worked.") - except Exception as e: # noqa: BLE001 + except Exception as e: _logger.debug("Error during tmp file removal.", exc_info=e) self._delete_manually(check_only=True) if not self.exists: diff --git a/tests/test_chunking.py b/tests/test_chunking.py new file mode 100644 index 00000000..cc905e41 --- /dev/null +++ b/tests/test_chunking.py @@ -0,0 +1,434 @@ +""" +Tests for sending commands that are too big for one write. + +The real limit is 100MB, which is too slow to test against, so nearly every +test here shrinks it with monkeypatch. The whole point of the feature is +that a chunked call and a normal call are indistinguishable, so most of +these run the same call both ways and compare. +""" + +import asyncio +import json + +import logistro +import pytest +import pytest_asyncio + +from choreographer.channels import MessageTooLargeError, pipe +from choreographer.protocol import _chunking + +pytestmark = pytest.mark.asyncio(loop_scope="function") + +_logger = logistro.getLogger(__name__) + +_N_VALUES = 2000 +_MIN_CHUNKED_MESSAGES = 3 # init + at least one push + the real call +_MIN_REAL_CHUNKED_MESSAGES = 5 # at the real 10MiB chunk size, expect about 14 + +_ECHO_FN = ( + "function(spec, tag){" + "return JSON.stringify({" + "type: typeof spec," + "n: spec.values.length," + "first: spec.values[0]," + "last: spec.values[spec.values.length - 1]," + "text: spec.text," + "tag: tag" + "});" + "}" +) + + +def _spec(n=_N_VALUES, text="plain"): + return {"values": list(range(n)), "text": text} + + +@pytest_asyncio.fixture(scope="function", loop_scope="function") +async def js(browser): + """Give back a tab session with a JavaScript context to run functions in.""" + tab = await browser.create_tab("") + session = await tab.create_session() + context = session.subscribe_once("Runtime.executionContextCreated") + await session.send_command("Page.enable") + await session.send_command("Runtime.enable") + js_id = (await context)["params"]["context"]["id"] + yield session, js_id + await tab.close_session(session.session_id) + await browser.close_tab(tab.target_id) + + +@pytest_asyncio.fixture(scope="function", loop_scope="function") +async def js_unique(browser): + """Like `js`, but hands back the context's uniqueId instead of its id.""" + tab = await browser.create_tab("") + session = await tab.create_session() + context = session.subscribe_once("Runtime.executionContextCreated") + await session.send_command("Page.enable") + await session.send_command("Runtime.enable") + unique_id = (await context)["params"]["context"]["uniqueId"] + yield session, unique_id + await tab.close_session(session.session_id) + await browser.close_tab(tab.target_id) + + +async def _call(session, js_id, fn=_ECHO_FN, args=None, **extra): + params = { + "functionDeclaration": fn, + "arguments": [{"value": a} for a in (args if args is not None else [])], + "executionContextId": js_id, + "returnByValue": False, + "awaitPromise": True, + } + params.update(extra) + return await session.send_command("Runtime.callFunctionOn", params=params) + + +def _value(response): + return response["result"]["result"]["value"] + + +async def _store_keys(session, js_id): + response = await _call( + session, + js_id, + fn=( + "function(){" + "return JSON.stringify(" + "window.__choreo_chunks ? Object.keys(window.__choreo_chunks) : []" + ");" + "}" + ), + ) + return json.loads(_value(response)) + + +def _shrink(monkeypatch, max_size=4096, chunk_size=512): + monkeypatch.setattr(pipe, "MAX_MESSAGE_SIZE", max_size) + monkeypatch.setattr(_chunking, "CHUNK_SIZE", chunk_size) + + +async def test_chunked_matches_unchunked(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + args = [_spec(), "hello"] + + plain = _value(await _call(session, js_id, args=args)) + + _shrink(monkeypatch) + chunked = _value(await _call(session, js_id, args=args)) + + assert json.loads(chunked) == json.loads(plain) + assert json.loads(chunked)["type"] == "object" + assert json.loads(chunked)["n"] == _N_VALUES + + +async def test_chunking_actually_happened(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + before = session.message_id + await _call(session, js_id, args=[_spec(), "hello"]) + assert session.message_id - before > _MIN_CHUNKED_MESSAGES + + +async def test_store_is_cleaned_up(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + await _call(session, js_id, args=[_spec(), "hello"]) + + monkeypatch.undo() + assert await _store_keys(session, js_id) == [] + + +async def test_error_in_function_still_reports_and_cleans_up(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + response = await _call( + session, + js_id, + fn="function(spec){ throw new Error('boom ' + spec.values.length); }", + args=[_spec()], + ) + assert "exceptionDetails" in response["result"] + assert "boom 2000" in json.dumps(response["result"]["exceptionDetails"]) + + monkeypatch.undo() + assert await _store_keys(session, js_id) == [] + + +async def test_promise_is_awaited(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + response = await _call( + session, + js_id, + fn=( + "function(spec){return Promise.resolve('resolved ' + spec.values.length);}" + ), + args=[_spec()], + ) + assert _value(response) == "resolved 2000" + + +async def test_return_by_value(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + response = await _call( + session, + js_id, + fn="function(spec){ return {n: spec.values.length}; }", + args=[_spec()], + returnByValue=True, + ) + assert _value(response) == {"n": 2000} + + +async def test_with_perf_survives_chunking(js, monkeypatch): + """ + `with_perf` has to keep working when the send gets broken up. + + Timings are looked up by message key, and the original command's write + never happened, so its key has no entry in `write_perfs`. `send_chunked` + hands back the command that actually went out so the lookup lands on that + one instead. Drop that and this raises `KeyError`. + + Note the timings then describe only the final message, not the pushes. + """ + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + response, perf = await session.send_command( + "Runtime.callFunctionOn", + params={ + "functionDeclaration": _ECHO_FN, + "arguments": [{"value": _spec()}, {"value": "hello"}], + "executionContextId": js_id, + "awaitPromise": True, + }, + with_perf=True, + ) + + assert json.loads(_value(response))["n"] == _N_VALUES + write_start, write_end, read_end = perf + assert write_start <= write_end <= read_end + + +async def test_non_ascii_survives_the_split(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + text = "héllo — 😀 中文 " * 200 + args = [_spec(text=text), "tag"] + + plain = json.loads(_value(await _call(session, js_id, args=args))) + + _shrink(monkeypatch) + chunked = json.loads(_value(await _call(session, js_id, args=args))) + + assert chunked["text"] == text + assert chunked == plain + + +async def test_several_large_arguments(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + response = await _call( + session, + js_id, + fn=( + "function(a, b, c){" + "return JSON.stringify([a.values.length, b.values.length, c]);" + "}" + ), + args=[_spec(1000), _spec(1500), "tail"], + ) + assert json.loads(_value(response)) == [1000, 1500, "tail"] + + +async def test_concurrent_chunked_calls_do_not_collide(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + responses = await asyncio.gather( + _call(session, js_id, args=[_spec(1000), "first"]), + _call(session, js_id, args=[_spec(2000), "second"]), + ) + first, second = (json.loads(_value(r)) for r in responses) + + assert (first["n"], first["tag"]) == (1000, "first") + assert (second["n"], second["tag"]) == (2000, "second") + + monkeypatch.undo() + assert await _store_keys(session, js_id) == [] + + +async def test_unchunkable_command_raises(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + with pytest.raises(MessageTooLargeError): + await session.send_command( + "Page.navigate", + params={"url": "data:text/html," + ("x" * 8192)}, + ) + + # The channel was never written to, so the session still works. + monkeypatch.undo() + assert _value(await _call(session, js_id, args=[_spec(10), "after"])) + + +async def test_lost_store_mid_send_raises(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + # Intentionally break init function to trigger this error + monkeypatch.setattr(_chunking, "_INIT_FN", "function(k){ }") + + with pytest.raises(RuntimeError, match="Chunked send failed in the page"): + await _call(session, js_id, args=[_spec(), "hello"]) + + +async def test_unique_context_id_is_followed(js_unique, monkeypatch): + _logger.info("testing...") + session, unique_id = js_unique + _shrink(monkeypatch) + + response = await session.send_command( + "Runtime.callFunctionOn", + params={ + "functionDeclaration": _ECHO_FN, + "arguments": [{"value": _spec()}, {"value": "hello"}], + "uniqueContextId": unique_id, + "awaitPromise": True, + }, + ) + + got = json.loads(_value(response)) + assert got["n"] == _N_VALUES + assert got["tag"] == "hello" + + +async def test_throw_on_side_effect_is_refused(js, monkeypatch): + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + with pytest.raises(MessageTooLargeError): + await _call(session, js_id, args=[_spec(), "hello"], throwOnSideEffect=True) + + # The important half: we declined before writing anything to the page. + monkeypatch.undo() + assert await _store_keys(session, js_id) == [] + + +async def test_object_id_arguments_are_not_chunked(js, monkeypatch): + """Browser-side handles can't be rebuilt from text, so don't try.""" + _logger.info("testing...") + session, js_id = js + handle = await _call( + session, + js_id, + fn="function(){ return {a: 1}; }", + ) + object_id = handle["result"]["result"]["objectId"] + + _shrink(monkeypatch) + with pytest.raises(MessageTooLargeError): + await session.send_command( + "Runtime.callFunctionOn", + params={ + "functionDeclaration": "function(o, pad){ return o.a; }", + "arguments": [{"objectId": object_id}, {"value": "x" * 8192}], + "executionContextId": js_id, + }, + ) + + +async def test_oversized_function_does_not_recurse(js, monkeypatch): + """ + A huge functionDeclaration can't be helped by chunking. + + The wrapper we'd build is bigger than the message we're replacing, so it + has to give up rather than fall back into itself forever. + """ + _logger.info("testing...") + session, js_id = js + _shrink(monkeypatch) + + with pytest.raises(MessageTooLargeError): + await _call( + session, + js_id, + fn="function(x){ var s = '" + ("y" * 8192) + "'; return s.length; }", + args=[1], + ) + + +@pytest.mark.slow +async def test_real_oversized_payload(js): + """ + The only test that uses Chrome's actual limit, so it is the slow one. + + The padding is stamped with its own offset every 1000 characters, so a + dropped piece or two pieces arriving out of order changes the answer. + A payload of all one character would hide both. + """ + _logger.info("testing...") + session, js_id = js + + stride = 1000 + n_pieces = 115_000 # About 110MiB (over the limit) + pad = "".join(f"{i:08d}" + "y" * (stride - 8) for i in range(n_pieces)) + marks = [0, 5_000, 50_000, n_pieces - 1] + spec = {"pad": pad, "tail": "héllo 😀"} + + before = session.message_id + response = await _call( + session, + js_id, + fn=( + "function(spec, offsets, stride){" + "return JSON.stringify({" + "len: spec.pad.length," + "tail: spec.tail," + "marks: offsets.map(function(k){" + "return spec.pad.substr(k * stride, 8);" + "})" + "});" + "}" + ), + args=[spec, marks, stride], + ) + got = json.loads(_value(response)) + sent = session.message_id - before + + # Without this the test would still pass if the payload ever slipped under + # the real limit, quietly measuring nothing. + assert sent > _MIN_REAL_CHUNKED_MESSAGES, ( + f"only {sent} messages sent: the payload did not get chunked" + ) + assert got["len"] == n_pieces * stride + assert got["tail"] == "héllo 😀" + assert got["marks"] == [f"{k:08d}" for k in marks] + assert await _store_keys(session, js_id) == [] + + +async def test_error_text_does_not_include_payload(): + _logger.info("testing...") + error = MessageTooLargeError(200, 100, payload="SECRET" * 100) + assert "SECRET" not in str(error) + assert "SECRET" not in repr(error) + assert error.payload is not None diff --git a/tests/test_serializer.py b/tests/test_serializer.py index a82af86f..214fb007 100644 --- a/tests/test_serializer.py +++ b/tests/test_serializer.py @@ -10,10 +10,11 @@ import json from typing import TYPE_CHECKING -import choreographer.channels._wire as wire import logistro import numpy as np import pytest + +import choreographer.channels._wire as wire from choreographer.channels import register_custom_encoder if TYPE_CHECKING: diff --git a/tests/test_session.py b/tests/test_session.py index f5dfc9e5..8c6e0213 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -3,6 +3,7 @@ import logistro import pytest import pytest_asyncio + from choreographer import errors # allows to create a browser pool for tests