zcbot/tests/test_file_store.py

"""原子文件写与知识库跨进程写锁回归。"""
from __future__ import annotations

import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path
from unittest.mock import patch
from uuid import UUID

from core.file_store import (
    FileLockBusy,
    atomic_write_text,
    interprocess_file_lock,
)
from core.kb import (
    create_kb,
    delete_doc,
    delete_kb,
    format_index_line,
    save_source,
    save_sources,
)
from core.kb_lock import KbBusyError, kb_is_locked, kb_mutation_lock
from tools.fs import WriteTool


_UID = UUID("11111111-2222-3333-4444-555555555555")


def _hold_lock_in_child(path: str, ready, release) -> None:
    with interprocess_file_lock(Path(path)):
        ready.set()
        release.wait(10)


class AtomicWriteTests(unittest.TestCase):
    def test_atomic_write_replaces_complete_file_and_cleans_temp(self):
        with tempfile.TemporaryDirectory() as td:
            path = Path(td) / "a.txt"
            path.write_text("old", encoding="utf-8")
            atomic_write_text(path, "新内容")
            self.assertEqual(path.read_text(encoding="utf-8"), "新内容")
            self.assertEqual(list(path.parent.glob(".a.txt.*.tmp")), [])

    def test_replace_failure_keeps_old_file_and_cleans_temp(self):
        with tempfile.TemporaryDirectory() as td:
            path = Path(td) / "a.txt"
            path.write_text("old", encoding="utf-8")
            with patch("core.file_store.os.replace", side_effect=OSError("boom")):
                with self.assertRaises(OSError):
                    atomic_write_text(path, "new")
            self.assertEqual(path.read_text(encoding="utf-8"), "old")
            self.assertEqual(list(path.parent.glob(".a.txt.*.tmp")), [])


class FileLockTests(unittest.TestCase):
    def test_second_handle_cannot_acquire_locked_file(self):
        with tempfile.TemporaryDirectory() as td:
            lock = Path(td) / "x.lock"
            with interprocess_file_lock(lock):
                with self.assertRaises(FileLockBusy):
                    with interprocess_file_lock(lock):
                        pass
            with interprocess_file_lock(lock):
                pass

    def test_lock_is_visible_across_processes(self):
        with tempfile.TemporaryDirectory() as td:
            lock = Path(td) / "cross-process.lock"
            ctx = get_context("spawn")
            ready = ctx.Event()
            release = ctx.Event()
            child = ctx.Process(
                target=_hold_lock_in_child, args=(str(lock), ready, release)
            )
            child.start()
            try:
                self.assertTrue(ready.wait(5), "child did not acquire lock")
                with self.assertRaises(FileLockBusy):
                    with interprocess_file_lock(lock):
                        pass
            finally:
                release.set()
                child.join(5)
                if child.is_alive():
                    child.terminate()
                    child.join(5)
            self.assertEqual(child.exitcode, 0)


class SandboxPackagingTests(unittest.TestCase):
    def test_sandbox_copies_file_store_dependencies(self):
        root = Path(__file__).resolve().parents[1]
        dockerfile = (root / "deploy" / "sandbox" / "Dockerfile").read_text("utf-8")
        self.assertIn("core/file_store.py", dockerfile)
        self.assertIn("core/kb_lock.py", dockerfile)


class KnowledgeBaseMutationTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.ws = Path(self.tmp.name)
        self.user_root = self.ws / "users" / str(_UID)
        self.kb = create_kb(self.ws, _UID, "标准库")
        assert self.kb is not None

    def tearDown(self):
        self.tmp.cleanup()

    def test_lock_blocks_service_and_agent_mutations(self):
        idx = self.kb / "INDEX.md"
        old = idx.read_text(encoding="utf-8")
        tool = WriteTool(base_dir=self.user_root, user_root=self.user_root)

        with kb_mutation_lock(self.ws, _UID, "标准库"):
            self.assertTrue(kb_is_locked(self.ws, _UID, "标准库"))
            with self.assertRaises(KbBusyError):
                delete_kb(self.ws, _UID, "标准库")
            result = tool.execute(".kb/标准库/INDEX.md", "bad")
            self.assertTrue(result.startswith("[Error]"))
            self.assertEqual(idx.read_text(encoding="utf-8"), old)

        self.assertFalse(kb_is_locked(self.ws, _UID, "标准库"))
        result = tool.execute(".kb/标准库/INDEX.md", "ok")
        self.assertTrue(result.startswith("[wrote"))
        self.assertEqual(idx.read_text(encoding="utf-8"), "ok")

    def test_invalid_batch_is_rejected_before_any_write(self):
        result = save_sources(
            self.ws,
            _UID,
            "标准库",
            [("good.txt", b"good"), ("../bad.txt", b"bad")],
        )
        self.assertIsNone(result)
        self.assertFalse((self.kb / "sources" / "good.txt").exists())

    def _seed_indexed_doc(self):
        source = self.kb / "sources" / "a.txt"
        doc = self.kb / "docs" / "a.md"
        source.write_bytes(b"old")
        doc.write_text("old doc", encoding="utf-8")
        line = format_index_line(
            title="A",
            doc="docs/a.md",
            source="sources/a.txt",
            summary="old",
            keywords="a",
        )
        (self.kb / "INDEX.md").write_text(
            f"# 标准库\n\n{line}\n", encoding="utf-8"
        )
        return source, doc

    def test_overwrite_removes_old_index_and_doc_then_publishes_source(self):
        source, doc = self._seed_indexed_doc()
        self.assertEqual(
            save_source(self.ws, _UID, "标准库", "a.txt", b"new"), "a.txt"
        )
        self.assertEqual(source.read_bytes(), b"new")
        self.assertFalse(doc.exists())
        self.assertNotIn("docs/a.md", (self.kb / "INDEX.md").read_text("utf-8"))

    def test_delete_keeps_data_when_index_publish_fails(self):
        source, doc = self._seed_indexed_doc()
        with patch("core.kb.atomic_write_text", side_effect=OSError("disk full")):
            self.assertFalse(delete_doc(self.ws, _UID, "标准库", "a.md"))
        self.assertTrue(source.exists())
        self.assertTrue(doc.exists())
        self.assertIn("docs/a.md", (self.kb / "INDEX.md").read_text("utf-8"))


if __name__ == "__main__":
    unittest.main()