[micro_wake_word] Include the local model file in bundles (#17604)

esphome/bundle.py

@@ -7,7 +7,7 @@ and compiled directly: ``esphome compile my_device.esphomebundle.tar.gz``
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import StrEnum
 import io
 import json
@@ -32,6 +32,8 @@ from esphome.core import CORE, EsphomeError
 
 _LOGGER = logging.getLogger(__name__)
 
+DOMAIN = "bundle"
+
 BUNDLE_EXTENSION = ".esphomebundle.tar.gz"
 MANIFEST_FILENAME = "manifest.json"
 CURRENT_MANIFEST_VERSION = 1
@@ -120,6 +122,32 @@ def _find_used_secret_keys(yaml_files: list[Path]) -> set[str]:
     return keys
 
 
+@dataclass
+class BundleData:
+    """Files components asked to include, keyed under DOMAIN in CORE.data."""
+
+    extra_files: list[Path] = field(default_factory=list)
+
+
+def _get_data() -> BundleData:
+    if DOMAIN not in CORE.data:
+        CORE.data[DOMAIN] = BundleData()
+    return CORE.data[DOMAIN]
+
+
+def add_bundle_file(path: Path) -> None:
+    """Register a file that a bundle must include.
+
+    Bundle discovery walks the validated config, so it only finds files the config
+    names. Components call this during validation for files it cannot see, such as a
+    file that is referenced from inside another file.
+
+    A relative path is taken as relative to the config directory. Files outside the
+    config directory are skipped when the bundle is built.
+    """
+    _get_data().extra_files.append(CORE.relative_config_path(path))
+
+
 @dataclass
 class BundleFile:
     """A file to include in the bundle."""
@@ -286,13 +314,18 @@ class ConfigBundleCreator:
         with known file extensions are also resolved and checked.
 
         Core ESPHome concepts that use relative paths or directories
-        are handled explicitly.
+        are handled explicitly. Files the config does not name at all are
+        registered by their component with add_bundle_file().
         """
         config = self._config
 
         # Generic walk: find all file paths in the validated config
         self._walk_config_for_files(config)
 
+        # Files registered by components during validation
+        for extra_file in _get_data().extra_files:
+            self._add_file(extra_file)
+
         # --- Core ESPHome concepts needing explicit handling ---
 
         # esphome.includes / includes_c - can be relative paths and directories

esphome/components/micro_wake_word/__init__.py

@@ -6,6 +6,7 @@ from urllib.parse import urljoin
 
 from esphome import automation, external_files, git
 from esphome.automation import register_action, register_condition
+from esphome.bundle import add_bundle_file
 import esphome.codegen as cg
 from esphome.components import esp32, microphone, ota, psram
 import esphome.config_validation as cv
@@ -28,6 +29,7 @@ from esphome.const import (
     TYPE_LOCAL,
 )
 from esphome.core import CORE, HexInt
+from esphome.types import ConfigType
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -236,10 +238,45 @@ HTTP_SCHEMA = cv.All(
     _process_http_source,
 )
 
-LOCAL_SCHEMA = cv.Schema(
-    {
-        cv.Required(CONF_PATH): cv.All(_validate_json_filename, cv.file_),
-    }
+
+def _register_local_model_file(config: ConfigType) -> ConfigType:
+    """Register the model file that the manifest points to, so bundles include it.
+
+    The manifest names its model file relative to itself, so that path never appears
+    in the YAML and bundle discovery cannot find it on its own.
+
+    Problems with the manifest are logged and ignored here rather than raised. Loading
+    the manifest later reports them with better messages, and raising would be
+    swallowed by the shorthand validator, which then reports a confusing error about a
+    missing file in a git repository. Logging keeps the skipped registration
+    diagnosable if the manifest is only briefly unreadable, since the bundle would
+    then be built without the model file.
+    """
+    manifest_path: Path = config[CONF_PATH]
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        model = manifest[CONF_MODEL]
+    except (OSError, ValueError, KeyError, TypeError) as err:
+        _LOGGER.debug("Not registering a model file from %s: %s", manifest_path, err)
+        return config
+    if not isinstance(model, str):
+        _LOGGER.debug(
+            "Not registering a model file from %s: 'model' is %s, expected a string",
+            manifest_path,
+            type(model).__name__,
+        )
+        return config
+    add_bundle_file(manifest_path.parent / model)
+    return config
+
+
+LOCAL_SCHEMA = cv.All(
+    cv.Schema(
+        {
+            cv.Required(CONF_PATH): cv.All(_validate_json_filename, cv.file_),
+        }
+    ),
+    _register_local_model_file,
 )
 
 

tests/component_tests/micro_wake_word/test_init.py

@@ -0,0 +1,110 @@
+"""Tests for micro_wake_word local model validation."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from esphome.components.micro_wake_word import LOCAL_SCHEMA
+from esphome.core import CORE
+
+MANIFEST: dict[str, Any] = {
+    "type": "micro",
+    "model": "hey_jarvis.tflite",
+    "author": "someone",
+    "version": 2,
+    "wake_word": "hey jarvis",
+    "trained_languages": ["en"],
+    "micro": {
+        "feature_step_size": 10,
+        "tensor_arena_size": 30000,
+        "probability_cutoff": 0.97,
+        "sliding_window_size": 5,
+        "minimum_esphome_version": "2024.7.0",
+    },
+}
+
+
+def _registered_files() -> list[Path]:
+    """Files components registered for bundling this run."""
+    data = CORE.data.get("bundle")
+    return list(data.extra_files) if data else []
+
+
+@pytest.fixture
+def config_dir(tmp_path: Path) -> Path:
+    """A config dir holding a manifest and its model file."""
+    (tmp_path / "models").mkdir()
+    (tmp_path / "models" / "hey_jarvis.tflite").write_bytes(b"fake model")
+    (tmp_path / "models" / "hey_jarvis.json").write_text(json.dumps(MANIFEST))
+    CORE.config_path = tmp_path / "test.yaml"
+    return tmp_path
+
+
+def test_local_schema_registers_model_file(config_dir: Path) -> None:
+    """The model file named by the manifest is registered so bundles include it."""
+    LOCAL_SCHEMA({"path": "models/hey_jarvis.json"})
+
+    assert _registered_files() == [config_dir / "models" / "hey_jarvis.tflite"]
+
+
+def test_local_schema_registers_model_file_in_subdirectory(config_dir: Path) -> None:
+    """The model reference is resolved relative to the manifest, not the config dir."""
+    nested = config_dir / "models" / "nested"
+    nested.mkdir()
+    (nested / "model.tflite").write_bytes(b"fake model")
+    (config_dir / "models" / "nested.json").write_text(
+        json.dumps({**MANIFEST, "model": "nested/model.tflite"})
+    )
+
+    LOCAL_SCHEMA({"path": "models/nested.json"})
+
+    assert _registered_files() == [nested / "model.tflite"]
+
+
+def test_local_schema_leaves_config_untouched(config_dir: Path) -> None:
+    """Registration is a side effect; the model file is not a config key."""
+    config = LOCAL_SCHEMA({"path": "models/hey_jarvis.json"})
+
+    assert config == {"path": config_dir / "models" / "hey_jarvis.json"}
+
+
+def test_local_schema_missing_model_file_still_validates(config_dir: Path) -> None:
+    """A model file that does not exist is registered, not rejected.
+
+    Raising here would be swallowed by the shorthand validator, which would then
+    report a confusing error about a missing file in a git repository.
+    """
+    (config_dir / "models" / "hey_jarvis.tflite").unlink()
+
+    LOCAL_SCHEMA({"path": "models/hey_jarvis.json"})
+
+    assert _registered_files() == [config_dir / "models" / "hey_jarvis.tflite"]
+
+
+@pytest.mark.parametrize(
+    "contents",
+    [
+        pytest.param("{not valid json", id="malformed"),
+        pytest.param(json.dumps({"type": "micro"}), id="no_model_key"),
+        pytest.param(json.dumps(["a", "list"]), id="not_an_object"),
+        pytest.param(json.dumps({"model": 42}), id="model_not_a_string"),
+    ],
+)
+def test_local_schema_bad_manifest_does_not_raise(
+    config_dir: Path, contents: str, caplog: pytest.LogCaptureFixture
+) -> None:
+    """Manifest problems are left to later stages, which report them better.
+
+    The skipped registration is logged so a bundle built without the model file can
+    be diagnosed.
+    """
+    (config_dir / "models" / "hey_jarvis.json").write_text(contents)
+
+    with caplog.at_level(logging.DEBUG):
+        LOCAL_SCHEMA({"path": "models/hey_jarvis.json"})
+
+    assert _registered_files() == []
+    assert "Not registering a model file" in caplog.text

tests/unit_tests/test_bundle.py

@@ -22,6 +22,7 @@ from esphome.bundle import (
     _add_bytes_to_tar,
     _default_target_dir,
     _find_used_secret_keys,
+    add_bundle_file,
     extract_bundle,
     is_bundle_path,
     prepare_bundle_for_compile,
@@ -611,6 +612,70 @@ def test_discover_files_includes_config(tmp_path: Path) -> None:
     assert "test.yaml" in paths
 
 
+def test_discover_files_includes_registered_files(tmp_path: Path) -> None:
+    """Files registered with add_bundle_file() are included.
+
+    The config does not name them, so discovery cannot find them on its own.
+    """
+    config_dir = _setup_config_dir(
+        tmp_path,
+        files={"models/model.tflite": "fake model data"},
+    )
+    add_bundle_file(config_dir / "models" / "model.tflite")
+
+    creator = ConfigBundleCreator({})
+    files = creator.discover_files()
+
+    paths = [f.path for f in files]
+    assert "models/model.tflite" in paths
+
+
+def test_discover_files_registered_relative_file(tmp_path: Path) -> None:
+    """A relative registered path is taken as relative to the config directory.
+
+    Not the working directory, which is where Path.resolve() would put it.
+    """
+    _setup_config_dir(
+        tmp_path,
+        files={"models/model.tflite": "fake model data"},
+    )
+    add_bundle_file(Path("models/model.tflite"))
+
+    creator = ConfigBundleCreator({})
+    files = creator.discover_files()
+
+    paths = [f.path for f in files]
+    assert "models/model.tflite" in paths
+
+
+def test_discover_files_registered_file_outside_config_dir(tmp_path: Path) -> None:
+    """A registered file outside the config directory is skipped, not bundled."""
+    _setup_config_dir(tmp_path)
+    outside = tmp_path / "outside.tflite"
+    outside.write_text("fake model data")
+    add_bundle_file(outside)
+
+    creator = ConfigBundleCreator({})
+    files = creator.discover_files()
+
+    assert [f.path for f in files] == ["test.yaml"]
+
+
+def test_discover_files_registered_file_deduplicated(tmp_path: Path) -> None:
+    """Registering the same file twice adds it once."""
+    config_dir = _setup_config_dir(
+        tmp_path,
+        files={"models/model.tflite": "fake model data"},
+    )
+    add_bundle_file(config_dir / "models" / "model.tflite")
+    add_bundle_file(config_dir / "models" / "model.tflite")
+
+    creator = ConfigBundleCreator({})
+    files = creator.discover_files()
+
+    assert [f.path for f in files].count("models/model.tflite") == 1
+
+
 def test_discover_files_finds_path_objects(tmp_path: Path) -> None:
     """Path objects in validated config are discovered."""
     config_dir = _setup_config_dir(