github/esphome
2
0
Fork 0
mirror of https://github.com/esphome/esphome.git synced 2026-09-30 08:20:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[store_yaml] Simplify: lazy YAML discovery in to_code, zero-copy streaming, shared envelope decoder

Browse Source
This commit is contained in:
J. Nick Koston
2026-07-03 19:38:41 -05:00
parent e4bb7c93a2
commit bd297f08a8
10 changed files with 160 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/test_store_yaml_recovery.py
+15 -35
View File
@@ -16,7 +16,6 @@ from __future__ import annotations
import asyncio
import contextlib
import struct
import pytest
@@ -25,6 +24,8 @@ try:
except ImportError:
from backports import zstd # type: ignore[import-not-found, no-redef]
from esphome.components.store_yaml import unpack_envelope
from .types import RunCompiledFunction
# Message IDs from esphome/components/api/api.proto.
@@ -33,8 +34,6 @@ HELLO_RESPONSE = 2
GET_YAML_REQUEST = 149
GET_YAML_RESPONSE = 150
ENVELOPE_MAGIC = b"EHY1"
def _encode_varint(value: int) -> bytes:
"""Encode an unsigned integer as a protobuf varint."""
@@ -94,25 +93,16 @@ def _parse_get_yaml_response(payload: bytes) -> tuple[bytes, bool, int, str]:
return data, done, total_size, encoding
def _unpack_envelope(blob: bytes) -> dict[str, bytes]:
"""Inverse of `_pack_envelope` in `esphome/components/store_yaml/__init__.py`."""
assert blob[:4] == ENVELOPE_MAGIC, "envelope must start with EHY1 magic"
pos = 4
(count,) = struct.unpack_from("<I", blob, pos)
pos += 4
files: dict[str, bytes] = {}
for _ in range(count):
(path_len,) = struct.unpack_from("<H", blob, pos)
pos += 2
path = blob[pos : pos + path_len].decode("utf-8")
pos += path_len
(content_len,) = struct.unpack_from("<I", blob, pos)
pos += 4
content = blob[pos : pos + content_len]
pos += content_len
files[path] = content
assert pos == len(blob), "envelope must consume all bytes"
return files
async def _read_varint_from(reader: asyncio.StreamReader) -> int:
"""Read a protobuf varint byte-by-byte from a stream."""
result = 0
shift = 0
while True:
byte = (await reader.readexactly(1))[0]
result |= (byte & 0x7F) << shift
if not (byte & 0x80):
return result
shift += 7
class _PlaintextClient:
@@ -137,18 +127,8 @@ class _PlaintextClient:
preamble = await self._reader.readexactly(1)
assert preamble == b"\x00", f"unexpected preamble {preamble!r}"
async def _read_varint_stream() -> int:
result = 0
shift = 0
while True:
byte = (await self._reader.readexactly(1))[0]
result |= (byte & 0x7F) << shift
if not (byte & 0x80):
return result
shift += 7
payload_size = await _read_varint_stream()
msg_type = await _read_varint_stream()
payload_size = await _read_varint_from(self._reader)
msg_type = await _read_varint_from(self._reader)
payload = await self._reader.readexactly(payload_size) if payload_size else b""
return msg_type, payload
@@ -215,7 +195,7 @@ async def test_store_yaml_recovery(
)
envelope = zstd.decompress(compressed)
files = _unpack_envelope(envelope)
files = unpack_envelope(envelope)
assert files, "envelope should contain at least one file"
combined = b"\n".join(files.values())
tests/unit_tests/components/test_store_yaml.py
+31 -55
View File
@@ -3,41 +3,19 @@
from __future__ import annotations
from pathlib import Path
import struct
import pytest
from esphome.components.store_yaml import (
ENVELOPE_MAGIC,
REDACTED_PLACEHOLDER,
_gather_files,
_pack_envelope,
unpack_envelope,
)
from esphome.core import CORE, EsphomeError
from esphome.yaml_util import DiscoveredYamlFiles
def _unpack_envelope(blob: bytes) -> dict[str, bytes]:
"""Inverse of `_pack_envelope` for assertions in tests."""
assert blob[:4] == ENVELOPE_MAGIC, "envelope must start with EHY1 magic"
pos = 4
(count,) = struct.unpack_from("<I", blob, pos)
pos += 4
files: dict[str, bytes] = {}
for _ in range(count):
(path_len,) = struct.unpack_from("<H", blob, pos)
pos += 2
path = blob[pos : pos + path_len].decode("utf-8")
pos += path_len
(content_len,) = struct.unpack_from("<I", blob, pos)
pos += 4
content = blob[pos : pos + content_len]
pos += content_len
files[path] = content
assert pos == len(blob), "envelope must consume all bytes"
return files
@pytest.fixture
def project(tmp_path: Path) -> Path:
"""Lay out a tiny ESPHome-like project: entry yaml, an include, and a secrets file."""
@@ -49,31 +27,24 @@ def project(tmp_path: Path) -> Path:
return project_dir
@pytest.fixture(autouse=True)
def _reset_core() -> None:
CORE.data.pop("yaml_sources", None)
CORE.config_path = None
yield
CORE.data.pop("yaml_sources", None)
CORE.config_path = None
def _set_sources(project_dir: Path, *names: str, secrets: tuple[str, ...] = ()) -> None:
def _sources(
project_dir: Path, *names: str, secrets: tuple[str, ...] = ()
) -> DiscoveredYamlFiles:
CORE.config_path = project_dir / "entry.yaml"
files = [project_dir / name for name in names]
secret_paths = {(project_dir / name).resolve() for name in secrets}
CORE.data["yaml_sources"] = DiscoveredYamlFiles(files, secret_paths)
return DiscoveredYamlFiles(files, secret_paths)
def test_gather_redacts_secrets_by_default(project: Path) -> None:
_set_sources(
discovered = _sources(
project,
"entry.yaml",
"wifi.yaml",
"secrets.yaml",
secrets=("secrets.yaml",),
)
files = dict(_gather_files(include_secrets=False))
files = dict(_gather_files(discovered, include_secrets=False))
assert files["secrets.yaml"] == REDACTED_PLACEHOLDER
assert b"SUPER_SECRET" not in files["secrets.yaml"]
assert files["wifi.yaml"] == (project / "wifi.yaml").read_bytes()
@@ -82,8 +53,10 @@ def test_gather_redacts_secrets_by_default(project: Path) -> None:
def test_gather_redacts_yml_extension(project: Path) -> None:
yml = project / "secrets.yml"
yml.write_text("api_key: OTHER_SECRET\n")
_set_sources(project, "entry.yaml", "secrets.yml", secrets=("secrets.yml",))
files = dict(_gather_files(include_secrets=False))
discovered = _sources(
project, "entry.yaml", "secrets.yml", secrets=("secrets.yml",)
)
files = dict(_gather_files(discovered, include_secrets=False))
assert files["secrets.yml"] == REDACTED_PLACEHOLDER
@@ -101,15 +74,17 @@ def test_gather_redacts_secret_symlinked_to_other_name(
# but stores the resolved path; mimic that here.
resolved = link.resolve()
CORE.config_path = project / "entry.yaml"
CORE.data["yaml_sources"] = DiscoveredYamlFiles([resolved], {resolved})
files = dict(_gather_files(include_secrets=False))
discovered = DiscoveredYamlFiles([resolved], {resolved})
files = dict(_gather_files(discovered, include_secrets=False))
assert REDACTED_PLACEHOLDER in files.values()
assert b"FROM_SYMLINK" not in b"".join(files.values())
def test_gather_embeds_secrets_when_opted_in(project: Path) -> None:
_set_sources(project, "entry.yaml", "secrets.yaml", secrets=("secrets.yaml",))
files = dict(_gather_files(include_secrets=True))
discovered = _sources(
project, "entry.yaml", "secrets.yaml", secrets=("secrets.yaml",)
)
files = dict(_gather_files(discovered, include_secrets=True))
assert b"SUPER_SECRET" in files["secrets.yaml"]
@@ -120,10 +95,8 @@ def test_gather_uses_relative_path_for_external_files(
sibling = tmp_path / "outside.yaml"
sibling.write_text("foo: bar\n")
CORE.config_path = project / "entry.yaml"
CORE.data["yaml_sources"] = DiscoveredYamlFiles(
[project / "entry.yaml", sibling], set()
)
files = dict(_gather_files(include_secrets=False))
discovered = DiscoveredYamlFiles([project / "entry.yaml", sibling], set())
files = dict(_gather_files(discovered, include_secrets=False))
# project root is `tmp_path/project`, sibling is in `tmp_path` so it
# resolves to `../outside.yaml`.
assert "../outside.yaml" in files
@@ -132,7 +105,7 @@ def test_gather_uses_relative_path_for_external_files(
def test_gather_raises_when_no_sources(project: Path) -> None:
CORE.config_path = project / "entry.yaml"
with pytest.raises(EsphomeError):
_gather_files(include_secrets=False)
_gather_files(DiscoveredYamlFiles(), include_secrets=False)
def test_gather_raises_on_unreadable_file(
@@ -140,7 +113,7 @@ def test_gather_raises_on_unreadable_file(
) -> None:
"""An unreadable tracked file fails the build instead of producing a
silently partial recovery blob."""
_set_sources(project, "entry.yaml", "wifi.yaml")
discovered = _sources(project, "entry.yaml", "wifi.yaml")
orig_read_bytes = Path.read_bytes
def fake_read_bytes(self: Path) -> bytes:
@@ -150,7 +123,7 @@ def test_gather_raises_on_unreadable_file(
monkeypatch.setattr(Path, "read_bytes", fake_read_bytes)
with pytest.raises(EsphomeError, match="wifi.yaml"):
_gather_files(include_secrets=False)
_gather_files(discovered, include_secrets=False)
def test_gather_warns_on_unresolved_includes(
@@ -159,11 +132,9 @@ def test_gather_warns_on_unresolved_includes(
"""Substitution-pathed includes that discovery could not capture produce a
warning naming them, so the user knows the blob is incomplete."""
CORE.config_path = project / "entry.yaml"
CORE.data["yaml_sources"] = DiscoveredYamlFiles(
[project / "entry.yaml"], set(), ["${board}.yaml"]
)
discovered = DiscoveredYamlFiles([project / "entry.yaml"], set(), ["${board}.yaml"])
with caplog.at_level("WARNING", logger="esphome.components.store_yaml"):
files = _gather_files(include_secrets=False)
files = _gather_files(discovered, include_secrets=False)
assert len(files) == 1
assert any(
"${board}.yaml" in r.message and "not contain" in r.message
@@ -177,16 +148,21 @@ def test_pack_envelope_roundtrip() -> None:
("wifi.yaml", b"ssid: a\n"),
]
blob = _pack_envelope(files)
assert _unpack_envelope(blob) == dict(files)
assert unpack_envelope(blob) == dict(files)
def test_pack_envelope_handles_utf8_paths() -> None:
files = [("dossiers/maison.yaml", b"foo: bar\n")]
blob = _pack_envelope(files)
assert _unpack_envelope(blob) == dict(files)
assert unpack_envelope(blob) == dict(files)
def test_pack_envelope_rejects_overlong_path() -> None:
long_path = "a" * (0xFFFF + 1)
with pytest.raises(EsphomeError):
_pack_envelope([(long_path, b"")])
def test_unpack_envelope_rejects_bad_magic() -> None:
with pytest.raises(EsphomeError):
unpack_envelope(b"NOPE" + b"\x00" * 4)
tests/unit_tests/test_config_normalization.py
-21
View File
@@ -151,27 +151,6 @@ def test_validate_config_warns_on_dropped_merge_key(
assert yaml_util.take_dropped_merge_keys() == []
@pytest.mark.parametrize("store_yaml_enabled", [False, True])
def test_load_config_gates_yaml_discovery_on_store_yaml(
tmp_path: Path, store_yaml_enabled: bool
) -> None:
"""YAML source discovery (a second full parse) only runs when a consumer
like store_yaml is configured."""
content = "esphome:\n name: test\nhost:\napi:\n"
if store_yaml_enabled:
content += "store_yaml:\n allow_unencrypted: true\n"
main = tmp_path / "main.yaml"
main.write_text(content)
CORE.config_path = main
result = config.load_config({})
assert not result.errors
assert ("yaml_sources" in CORE.data) == store_yaml_enabled
if store_yaml_enabled:
assert main.resolve() in CORE.data["yaml_sources"].files
def test_validate_config_suppresses_merge_warning(
tmp_path: Path, caplog: pytest.LogCaptureFixture
) -> None:
tests/unit_tests/test_yaml_util.py
+1 -1
View File
@@ -1111,7 +1111,7 @@ def test_force_load_include_files_warns_on_load_failure(
def test_discovered_yaml_files_holds_files_and_secrets() -> None:
"""`DiscoveredYamlFiles` is a small data carrier; both fields are mandatory."""
"""`DiscoveredYamlFiles` is a small data carrier."""
files = [Path("/tmp/a.yaml")]
secrets = {Path("/tmp/a.yaml")}
discovered = DiscoveredYamlFiles(files, secrets)