Register a value written to a secrets file, keep the registry through own_documents, finish the restore after any failure and raise an interrupt again with the report

esphome/yaml_edit.py

@@ -106,7 +106,7 @@ def _own_documents() -> set[Path]:
     walk(CORE.raw_config or {})
     # A file that only gave substitutions leaves no node behind; the tracker
     # also lists a secrets file the loader only tried, hence the is_file check
-    found = yaml_util.discover_user_yaml_files(CORE.config_path)
+    found = yaml_util.discover_user_yaml_files(CORE.config_path, clear_secrets=False)
     if found.load_errors:  # an incomplete set would misplace a key
         raise EsphomeError(
             "Could not read every file of the configuration: "
@@ -118,7 +118,8 @@ def _own_documents() -> set[Path]:
 
 def source_of(mapping: ConfigType, name: str) -> tuple[Path, int] | None:
     """The file and line ``name:`` was read from; None when validation added
-    it or a merge key brought it in from an anchor elsewhere."""
+    it or a merge key brought it in from an anchor elsewhere. A mapping
+    merge_config rebuilt is placed by its winning value's line."""
     rng = getattr(next((k for k in mapping if k == name), None), "esp_range", None)
     if rng is None:
         return None
@@ -224,7 +225,10 @@ def apply_line_edits(edits: list[LineEdit]) -> dict[Path, Snapshot]:
     try:
         for path, snapshot in originals.items():
             write_keeping_mode(path, snapshot.written)
-        for path in originals:
+        reloads = list(originals)
+        if (main := CORE.config_path.resolve()) not in originals:
+            reloads.append(main)  # its !secret tags register a value written elsewhere
+        for path in reloads:
             try:
                 # The secret registry stays: later output still redacts
                 yaml_util.load_yaml(
@@ -247,9 +251,19 @@ def apply_line_edits(edits: list[LineEdit]) -> dict[Path, Snapshot]:
 
 def restore_files(originals: dict[Path, Snapshot]) -> None:
     """Put back what apply_line_edits wrote; a file that failed or changed
-    meanwhile is left alone and reported. An interrupt is reported the same
-    way, so the other files are still put back."""
-    failed = []
+    meanwhile is left alone and reported. Every file is tried whatever went
+    wrong before; an interrupt is raised again at the end with the report
+    as a note."""
+    failed: list[str] = []
+    stop: BaseException | None = None
+
+    def report(what: str, err: BaseException) -> None:
+        nonlocal stop
+        why = "interrupted" if isinstance(err, KeyboardInterrupt) else str(err)
+        failed.append(f"{what}: {why or type(err).__name__}")
+        if not isinstance(err, Exception):
+            stop = err
+
     for path, snapshot in originals.items():
         try:
             current = read_text(path)
@@ -258,11 +272,16 @@ def restore_files(originals: dict[Path, Snapshot]) -> None:
             if current != snapshot.written:
                 raise EsphomeError("changed since it was written, left as is")
             write_keeping_mode(path, snapshot.original)
-        except (EsphomeError, KeyboardInterrupt) as err:
-            failed.append(f"{path}: {str(err) or 'interrupted'}")
+        except (EsphomeError, OSError, KeyboardInterrupt) as err:
+            report(str(path), err)
     try:
         compiled_config.invalidate_compiled_config()
-    except EsphomeError as err:
-        failed.append(str(err))
-    if failed:
-        raise RestoreError("Could not restore " + "; ".join(failed))
+    except (EsphomeError, KeyboardInterrupt) as err:
+        report("the validated caches", err)
+    if not failed:
+        return
+    message = "Could not restore " + "; ".join(failed)
+    if stop is not None:
+        stop.add_note(message)
+        raise stop
+    raise RestoreError(message)

esphome/yaml_util.py

@@ -637,9 +637,13 @@ class DiscoveredYamlFiles:
     load_errors: list[str] = field(default_factory=list)
 
 
-def discover_user_yaml_files(config_path: Path) -> DiscoveredYamlFiles:
+def discover_user_yaml_files(
+    config_path: Path, clear_secrets: bool = True
+) -> DiscoveredYamlFiles:
     """Fresh-re-parse ``config_path`` and report every file the YAML loader
     pulled in, plus which of them came in under a secrets filename.
+    ``clear_secrets`` is passed to the load; False keeps the redaction
+    registry of a configuration loaded earlier in the run.
 
     Does NOT run schema validation, substitutions, or package resolution — so
     component-internal YAML loaded by validators (LVGL helpers, dashboard
@@ -663,7 +667,7 @@ def discover_user_yaml_files(config_path: Path) -> DiscoveredYamlFiles:
         _load_listeners.append(_capture_secret)
         try:
             try:
-                data = load_yaml(config_path)
+                data = load_yaml(config_path, clear_secrets=clear_secrets)
             except EsphomeError as err:
                 return DiscoveredYamlFiles(
                     list(loaded), secrets, load_errors=[f"{config_path}: {err}"]

tests/unit_tests/test_yaml_edit.py

@@ -335,7 +335,9 @@ def test_restore_reports_every_file_it_could_not_write(tmp_path: Path) -> None:
             "esphome.compiled_config.invalidate_compiled_config",
             side_effect=EsphomeError("busy"),
         ),
-        pytest.raises(EsphomeError, match="Could not restore busy"),
+        pytest.raises(
+            EsphomeError, match="Could not restore the validated caches: busy"
+        ),
     ):
         restore_files({path: Snapshot(YAML, YAML)})
     originals = apply_line_edits([_name_edit("garage")])
@@ -432,16 +434,65 @@ def test_restore_goes_on_after_an_interrupt(tmp_path: Path) -> None:
     )
     with (
         patch("esphome.yaml_edit.write_keeping_mode", side_effect=interrupt_at(1)),
-        pytest.raises(EsphomeError, match="Could not restore .*: interrupted"),
+        pytest.raises(KeyboardInterrupt) as info,
     ):
         restore_files(originals)
     assert calls == [path, other] and other.read_text() == "x: 1\n"
-    # Inside apply the failed rollback is reported with the reason it ran
+    assert f"Could not restore {path}: interrupted" in "".join(info.value.__notes__)
+    # The interrupt stays one inside apply as well
     _setup(tmp_path, YAML)
     calls.clear()
     with (
         patch("esphome.yaml_util.load_yaml", side_effect=EsphomeError("broken")),
         patch("esphome.yaml_edit.write_keeping_mode", side_effect=interrupt_at(2)),
-        pytest.raises(RestoreError, match="broken; Could not restore .*interrupted"),
+        pytest.raises(KeyboardInterrupt) as info,
     ):
         apply_line_edits([_name_edit("garage")])
+    assert "Could not restore" in "".join(info.value.__notes__)
+
+
+def test_restore_finishes_after_any_failure(tmp_path: Path) -> None:
+    """A raw OSError still leaves the other files put back, and an interrupt
+    while dropping the caches carries the report."""
+    path = _setup(tmp_path, YAML)
+    other = tmp_path / "other.yaml"
+    other.write_bytes(b"x: 1\n")
+    originals = apply_line_edits(
+        [_name_edit("garage"), LineEdit(other, 0, "x: 1", "x: 2")]
+    )
+    writes: list[Path] = []
+
+    def fail_first(target: Path, text: str) -> None:
+        writes.append(target)
+        if len(writes) == 1:
+            raise OSError("odd")
+        write_keeping_mode(target, text)
+
+    with (
+        patch("esphome.yaml_edit.write_keeping_mode", side_effect=fail_first),
+        pytest.raises(RestoreError, match=f"Could not restore {path}: odd"),
+    ):
+        restore_files(originals)
+    assert "garage" in path.read_text() and other.read_text() == "x: 1\n"
+    with (
+        patch(
+            "esphome.compiled_config.invalidate_compiled_config",
+            side_effect=KeyboardInterrupt,
+        ),
+        pytest.raises(KeyboardInterrupt) as info,
+    ):
+        restore_files(originals)
+    assert path.read_text() == YAML
+    assert "the validated caches: interrupted" in "".join(info.value.__notes__)
+
+
+def test_a_value_written_to_the_secrets_file_is_registered(tmp_path: Path) -> None:
+    """Only a resolved !secret registers its value, so the main file is
+    loaded again after an edit elsewhere; own_documents keeps the registry."""
+    secrets = tmp_path / "secrets.yaml"
+    secrets.write_bytes(b"ssid: hidden\n")
+    _setup(tmp_path, YAML.replace("ssid: kitchen", "ssid: !secret ssid"))
+    own_documents()
+    assert yaml_util.is_secret("hidden") == "ssid"
+    apply_line_edits([LineEdit(secrets, 0, "ssid: hidden", "ssid: hidden2")])
+    assert yaml_util.is_secret("hidden2") == "ssid"