From b117bbc12534e26db8aecb8d90e9c88394b0e1b5 Mon Sep 17 00:00:00 2001
From: teknium1 <teknium1@gmail.com>
Date: Sat, 14 Mar 2026 22:31:51 -0700
Subject: [PATCH] test: cover atomic temp cleanup on interrupts

- add regression coverage for BaseException cleanup in atomic_json_write
- add dedicated atomic_yaml_write tests, including interrupt cleanup
- document why BaseException is intentional in both helpers
---
 tests/test_atomic_json_write.py | 16 ++++++++++++
 tests/test_atomic_yaml_write.py | 44 +++++++++++++++++++++++++++++++++
 utils.py                        |  4 +++
 3 files changed, 64 insertions(+)
 create mode 100644 tests/test_atomic_yaml_write.py

diff --git a/tests/test_atomic_json_write.py b/tests/test_atomic_json_write.py
index cb8b2d6d..08bed89f 100644
--- a/tests/test_atomic_json_write.py
+++ b/tests/test_atomic_json_write.py
@@ -68,6 +68,22 @@ class TestAtomicJsonWrite:
         tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name]
         assert len(tmp_files) == 0
 
+    def test_cleans_up_temp_file_on_baseexception(self, tmp_path):
+        class SimulatedAbort(BaseException):
+            pass
+
+        target = tmp_path / "data.json"
+        original = {"preserved": True}
+        target.write_text(json.dumps(original), encoding="utf-8")
+
+        with patch("utils.json.dump", side_effect=SimulatedAbort):
+            with pytest.raises(SimulatedAbort):
+                atomic_json_write(target, {"new": True})
+
+        tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name]
+        assert len(tmp_files) == 0
+        assert json.loads(target.read_text(encoding="utf-8")) == original
+
     def test_accepts_string_path(self, tmp_path):
         target = str(tmp_path / "string_path.json")
         atomic_json_write(target, {"string": True})
diff --git a/tests/test_atomic_yaml_write.py b/tests/test_atomic_yaml_write.py
new file mode 100644
index 00000000..6a9e4f00
--- /dev/null
+++ b/tests/test_atomic_yaml_write.py
@@ -0,0 +1,44 @@
+"""Tests for utils.atomic_yaml_write — crash-safe YAML file writes."""
+
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+import yaml
+
+from utils import atomic_yaml_write
+
+
+class TestAtomicYamlWrite:
+    def test_writes_valid_yaml(self, tmp_path):
+        target = tmp_path / "data.yaml"
+        data = {"key": "value", "nested": {"a": 1}}
+
+        atomic_yaml_write(target, data)
+
+        assert yaml.safe_load(target.read_text(encoding="utf-8")) == data
+
+    def test_cleans_up_temp_file_on_baseexception(self, tmp_path):
+        class SimulatedAbort(BaseException):
+            pass
+
+        target = tmp_path / "data.yaml"
+        original = {"preserved": True}
+        target.write_text(yaml.safe_dump(original), encoding="utf-8")
+
+        with patch("utils.yaml.dump", side_effect=SimulatedAbort):
+            with pytest.raises(SimulatedAbort):
+                atomic_yaml_write(target, {"new": True})
+
+        tmp_files = [f for f in tmp_path.iterdir() if ".tmp" in f.name]
+        assert len(tmp_files) == 0
+        assert yaml.safe_load(target.read_text(encoding="utf-8")) == original
+
+    def test_appends_extra_content(self, tmp_path):
+        target = tmp_path / "data.yaml"
+
+        atomic_yaml_write(target, {"key": "value"}, extra_content="\n# comment\n")
+
+        text = target.read_text(encoding="utf-8")
+        assert "key: value" in text
+        assert "# comment" in text
diff --git a/utils.py b/utils.py
index 762bcb84..66d55290 100644
--- a/utils.py
+++ b/utils.py
@@ -50,6 +50,8 @@ def atomic_json_write(
             os.fsync(f.fileno())
         os.replace(tmp_path, path)
     except BaseException:
+        # Intentionally catch BaseException so temp-file cleanup still runs for
+        # KeyboardInterrupt/SystemExit before re-raising the original signal.
         try:
             os.unlink(tmp_path)
         except OSError:
@@ -96,6 +98,8 @@ def atomic_yaml_write(
             os.fsync(f.fileno())
         os.replace(tmp_path, path)
     except BaseException:
+        # Match atomic_json_write: cleanup must also happen for process-level
+        # interruptions before we re-raise them.
         try:
             os.unlink(tmp_path)
         except OSError: