203 lines
7.6 KiB
Python
203 lines
7.6 KiB
Python
"""Unit tests for atomic file-write utilities (U3 — R6).
|
|
|
|
Covers:
|
|
- Happy path: write_text_atomic writes correct content
|
|
- Edge case: no temp file residue after success
|
|
- Edge case: original file intact when write fails (mocked os.replace)
|
|
- Edge case: concurrent writes don't produce mixed content
|
|
- Error path: permission error on target directory
|
|
- Integration: settings _write_yaml_config uses atomic write
|
|
- Integration: settings _write_env_var uses atomic write
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from agentkit.server.utils.atomic_write import write_text_atomic
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Happy path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWriteTextAtomic:
|
|
def test_writes_content_correctly(self, tmp_path: Path):
|
|
target = tmp_path / "output.yaml"
|
|
write_text_atomic(target, "key: value\n")
|
|
assert target.read_text(encoding="utf-8") == "key: value\n"
|
|
|
|
def test_overwrites_existing_file(self, tmp_path: Path):
|
|
target = tmp_path / "output.yaml"
|
|
target.write_text("old: content\n", encoding="utf-8")
|
|
write_text_atomic(target, "new: content\n")
|
|
assert target.read_text(encoding="utf-8") == "new: content\n"
|
|
|
|
def test_no_temp_file_residue(self, tmp_path: Path):
|
|
target = tmp_path / "output.yaml"
|
|
write_text_atomic(target, "key: value\n")
|
|
# No temp/partial files should remain in the directory.
|
|
residue = [
|
|
f.name
|
|
for f in tmp_path.iterdir()
|
|
if f.name != "output.yaml" and not f.name.startswith(".")
|
|
]
|
|
assert residue == [], f"Unexpected temp files: {residue}"
|
|
|
|
def test_accepts_str_path(self, tmp_path: Path):
|
|
target = str(tmp_path / "output.txt")
|
|
write_text_atomic(target, "hello\n")
|
|
assert Path(target).read_text(encoding="utf-8") == "hello\n"
|
|
|
|
def test_creates_parent_dirs(self, tmp_path: Path):
|
|
target = tmp_path / "subdir" / "nested" / "output.yaml"
|
|
write_text_atomic(target, "key: value\n")
|
|
assert target.read_text(encoding="utf-8") == "key: value\n"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Crash safety
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWriteTextAtomicCrashSafety:
|
|
def test_original_intact_when_replace_fails(self, tmp_path: Path):
|
|
"""If os.replace raises, the original file must be unchanged."""
|
|
target = tmp_path / "output.yaml"
|
|
target.write_text("original: true\n", encoding="utf-8")
|
|
|
|
with patch("agentkit.server.utils.atomic_write.os.replace") as mock_replace:
|
|
mock_replace.side_effect = OSError("simulated crash")
|
|
with pytest.raises(OSError, match="simulated crash"):
|
|
write_text_atomic(target, "new: true\n")
|
|
|
|
# Original content must be intact.
|
|
assert target.read_text(encoding="utf-8") == "original: true\n"
|
|
# Temp file must be cleaned up.
|
|
residue = [
|
|
f.name
|
|
for f in tmp_path.iterdir()
|
|
if f.name != "output.yaml" and not f.name.startswith(".")
|
|
]
|
|
assert residue == [], f"Temp file not cleaned up: {residue}"
|
|
|
|
def test_original_intact_when_fsync_fails(self, tmp_path: Path):
|
|
"""If fsync raises (e.g. disk full), the original must be intact."""
|
|
target = tmp_path / "output.yaml"
|
|
target.write_text("original: true\n", encoding="utf-8")
|
|
|
|
with patch("agentkit.server.utils.atomic_write.os.fsync") as mock_fsync:
|
|
mock_fsync.side_effect = OSError("disk full")
|
|
with pytest.raises(OSError, match="disk full"):
|
|
write_text_atomic(target, "new: true\n")
|
|
|
|
assert target.read_text(encoding="utf-8") == "original: true\n"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Concurrency
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWriteTextAtomicConcurrency:
|
|
def test_concurrent_writes_no_mixed_content(self, tmp_path: Path):
|
|
"""Two threads writing different content — no mixed lines."""
|
|
target = tmp_path / "output.yaml"
|
|
target.write_text("init: 0\n", encoding="utf-8")
|
|
|
|
barrier = threading.Barrier(2)
|
|
results: list[bool] = []
|
|
|
|
def writer(content: str) -> None:
|
|
barrier.wait()
|
|
write_text_atomic(target, content)
|
|
results.append(True)
|
|
|
|
t1 = threading.Thread(target=writer, args=("aaa: 1\nbbb: 2\nccc: 3\n",))
|
|
t2 = threading.Thread(target=writer, args=("xxx: 1\nyyy: 2\nzzz: 3\n",))
|
|
t1.start()
|
|
t2.start()
|
|
t1.join(timeout=5)
|
|
t2.join(timeout=5)
|
|
|
|
assert len(results) == 2
|
|
# The final content must be one of the two writes, not a mix.
|
|
final = target.read_text(encoding="utf-8")
|
|
assert final in ("aaa: 1\nbbb: 2\nccc: 3\n", "xxx: 1\nyyy: 2\nzzz: 3\n")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Error paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWriteTextAtomicErrors:
|
|
def test_permission_error_on_unwritable_dir(self, tmp_path: Path):
|
|
"""Writing to a read-only directory raises PermissionError."""
|
|
ro_dir = tmp_path / "readonly"
|
|
ro_dir.mkdir()
|
|
os.chmod(ro_dir, 0o444)
|
|
try:
|
|
target = ro_dir / "output.yaml"
|
|
with pytest.raises(PermissionError):
|
|
write_text_atomic(target, "key: value\n")
|
|
finally:
|
|
os.chmod(ro_dir, 0o755)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration: settings.py
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSettingsAtomicWrite:
|
|
"""Verify that settings.py helpers use atomic writes."""
|
|
|
|
def test_write_env_var_uses_atomic_write(self, tmp_path: Path):
|
|
"""_write_env_var should write .env atomically (no partial writes)."""
|
|
from agentkit.server.routes.settings import _write_env_var
|
|
|
|
config_path = tmp_path / "agentkit.yaml"
|
|
config_path.write_text("llm:\n default_model: gpt-4\n", encoding="utf-8")
|
|
|
|
_write_env_var(str(config_path), "OPENAI_API_KEY", "sk-test-123")
|
|
|
|
env_path = tmp_path / ".env"
|
|
assert env_path.exists()
|
|
content = env_path.read_text(encoding="utf-8")
|
|
assert "OPENAI_API_KEY=sk-test-123" in content
|
|
|
|
# No temp files should remain.
|
|
residue = [
|
|
f.name
|
|
for f in tmp_path.iterdir()
|
|
if not f.name.startswith(".") and f.name != "agentkit.yaml"
|
|
]
|
|
assert residue == [], f"Unexpected files: {residue}"
|
|
|
|
def test_write_env_var_preserves_existing_content(self, tmp_path: Path):
|
|
"""Updating an existing key preserves other lines."""
|
|
from agentkit.server.routes.settings import _write_env_var
|
|
|
|
config_path = tmp_path / "agentkit.yaml"
|
|
config_path.write_text("llm:\n default_model: gpt-4\n", encoding="utf-8")
|
|
env_path = tmp_path / ".env"
|
|
env_path.write_text(
|
|
"# Header comment\nOPENAI_API_KEY=old-value\nANTHROPIC_API_KEY=sk-ant\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
_write_env_var(str(config_path), "OPENAI_API_KEY", "new-value")
|
|
|
|
content = env_path.read_text(encoding="utf-8")
|
|
assert "# Header comment" in content
|
|
assert "OPENAI_API_KEY=new-value" in content
|
|
assert "ANTHROPIC_API_KEY=sk-ant" in content
|
|
assert "old-value" not in content
|