diff --git a/tests/integration/fixtures/store_yaml_recovery_backpressure.yaml b/tests/integration/fixtures/store_yaml_recovery_backpressure.yaml new file mode 100644 index 00000000000..e2b0861a485 --- /dev/null +++ b/tests/integration/fixtures/store_yaml_recovery_backpressure.yaml @@ -0,0 +1,24 @@ +esphome: + name: store-yaml-backpressure + +host: + +logger: + +api: + +external_components: + - source: + type: local + path: EXTERNAL_COMPONENT_PATH + components: [sndbuf_pin_component] + +# Pins the device's socket send buffers for deterministic TCP backpressure +sndbuf_pin_component: + buffer_size: SERVER_SNDBUF + +store_yaml: + allow_unencrypted: true + +substitutions: +# GENERATED_SUBSTITUTIONS diff --git a/tests/integration/test_store_yaml_recovery.py b/tests/integration/test_store_yaml_recovery.py index b07f76372dd..dd300d407a9 100644 --- a/tests/integration/test_store_yaml_recovery.py +++ b/tests/integration/test_store_yaml_recovery.py @@ -5,6 +5,9 @@ Uses the raw client: released aioesphomeapi does not know GetYamlResponse yet. from __future__ import annotations +import asyncio +import random + from aioesphomeapi import api_pb2 import pytest @@ -60,3 +63,46 @@ async def test_store_yaml_recovery( assert b"recoverme123" not in envelope assert "ota_password" in find_secret_references(combined.decode()) assert b"ota_password: ''" in files["secrets.yaml"] + + +# Random hex barely compresses, so the blob spans many chunks and outgrows both buffers +SERVER_SNDBUF = 8192 +RECV_BUFFER_SIZE = 4096 +CHUNK_LIMIT = 1390 + + +@pytest.mark.asyncio +async def test_store_yaml_recovery_backpressure( + yaml_config: str, + run_compiled: RunCompiledFunction, + unused_tcp_port: int, +) -> None: + """A client that stops reading mid-stream still recovers every byte once it drains.""" + rng = random.Random(16445) + values = {f"blob_{i}": rng.randbytes(60).hex() for i in range(1000)} + config = yaml_config.replace("SERVER_SNDBUF", str(SERVER_SNDBUF)).replace( + "# GENERATED_SUBSTITUTIONS", + "\n".join(f" {name}: {value}" for name, value in values.items()), + ) + async with ( + run_compiled(config), + RawApiClient(unused_tcp_port, recv_buffer_size=RECV_BUFFER_SIZE) as client, + ): + await client.connect("store_yaml backpressure test") + await client.send_raw(GET_YAML_REQUEST, b"") + # Stop reading so the device fills both buffers and has to resume later + await asyncio.sleep(0.5) + responses: list[dict[int, list[int | bytes]]] = [] + while not responses or responses[-1].get(DONE_FIELD) != [1]: + responses.append(decode_fields(await client.read_frame(GET_YAML_RESPONSE))) + + chunks = [r[DATA_FIELD][0] for r in responses] + compressed = b"".join(chunks) + assert len(compressed) > 4 * (SERVER_SNDBUF + RECV_BUFFER_SIZE) + assert len(responses) > 1 and all(len(c) <= CHUNK_LIMIT for c in chunks) + assert all(r.get(DONE_FIELD) != [1] for r in responses[:-1]) + assert responses[0][TOTAL_SIZE_FIELD] == [len(compressed)] + combined = b"\n".join( + unpack_envelope(zstd_module().decompress(compressed)).values() + ) + assert all(value.encode() in combined for value in values.values()) diff --git a/tests/unit_tests/components/test_store_yaml.py b/tests/unit_tests/components/test_store_yaml.py index 83df45fe90b..ca6101e3017 100644 --- a/tests/unit_tests/components/test_store_yaml.py +++ b/tests/unit_tests/components/test_store_yaml.py @@ -1219,7 +1219,7 @@ def test_redacted_accepts_include_text_that_is_not_a_tag( def test_credential_words_cover_every_sensitive_key() -> None: """A removed branch is matched by name, so every `cv.sensitive` key must be named.""" root = Path(store_yaml.__file__).parents[2] - sources = [f.read_text() for f in root.rglob("*.py")] + sources = [f.read_text(encoding="utf-8") for f in root.rglob("*.py")] consts = { m.group(1): m.group(2) for text in sources