Harden backup import and restore coverage

1 files changed, 120 insertions, 0 deletions

diff --git a/tests/import_export_test.py b/tests/import_export_test.py
new file mode 100644
index 0000000..0ba5c5b
--- /dev/null
+++ b/tests/import_export_test.py
@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import tempfile
+from pathlib import Path
+
+
+MESSAGE_ID = "01ARZ3NDEKTSV4RRFFQ69G5FAV"
+
+
+def record(content: str = "valid", **changes: object) -> dict[str, object]:
+    value: dict[str, object] = {
+        "message_id": MESSAGE_ID,
+        "session_id": "escape-session",
+        "timestamp": 0,
+        "updated_at": 9223372036854775807,
+        "revision": 2147483647,
+        "role": "user",
+        "content": content,
+        "metadata": {
+            "tool_call_id": "tool\\\"\n",
+            "tags": ["tag\\\"\n"],
+            "memory_tier": "short",
+            "memory_kind": "conversation",
+            "importance": 100,
+            "summary": "summary\\\"\n\u3042",
+            "relations": [],
+            "source_message_ids": [],
+        },
+    }
+    value.update(changes)
+    return value
+
+
+def write_jsonl(path: Path, item: dict[str, object], final_newline: bool = True) -> None:
+    text = json.dumps({"type": "tiny-convmem-manifest", "format_version": 2}) + "\n"
+    text += json.dumps(item, ensure_ascii=True, separators=(",", ":"))
+    path.write_text(text + ("\n" if final_newline else ""), encoding="utf-8")
+
+
+def run(server: Path, data: Path, *arguments: str, success: bool) -> subprocess.CompletedProcess[str]:
+    result = subprocess.run(
+        [str(server), "--data-dir", str(data), *arguments],
+        capture_output=True,
+        text=True,
+    )
+    assert (result.returncode == 0) is success, (
+        result.returncode,
+        result.args,
+        result.stdout,
+        result.stderr,
+    )
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--server", required=True, type=Path)
+    args = parser.parse_args()
+    server = args.server.resolve()
+    root = Path(tempfile.mkdtemp(prefix="tiny-convmem-import-e2e-"))
+
+    prefix = "prefix\nquote=\" slash=\\ unicode=\u3042 "
+    escaped = prefix + "x" * (1024 * 1024 - len(prefix.encode("utf-8")))
+    assert len(escaped.encode("utf-8")) == 1024 * 1024
+    source = root / "source.jsonl"
+    write_jsonl(source, record(escaped))
+    first = root / "first"
+    run(server, first, "--import", str(source), success=True)
+    exported = root / "export.jsonl"
+    run(server, first, "--export", str(exported), success=True)
+    second = root / "second"
+    run(server, second, "--import", str(exported), success=True)
+    exported_again = root / "export-again.jsonl"
+    run(server, second, "--export", str(exported_again), success=True)
+    assert exported.read_bytes() == exported_again.read_bytes()
+    assert json.loads(exported.read_text(encoding="utf-8").splitlines()[1])["content"] == escaped
+
+    invalid_values: list[tuple[str, object]] = [
+        ("timestamp", -1),
+        ("timestamp", 9223372036854775808),
+        ("timestamp", 1.5),
+        ("timestamp", 1e3),
+        ("revision", 0),
+        ("revision", 2147483648),
+    ]
+    for index, (field, value) in enumerate(invalid_values):
+        invalid = root / f"invalid-{index}.jsonl"
+        write_jsonl(invalid, record(**{field: value}))
+        run(server, root / f"invalid-data-{index}", "--import", str(invalid), success=False)
+
+    malformed = root / "embedded-nul.jsonl"
+    write_jsonl(malformed, record("before\0after"))
+    run(server, root / "nul-data", "--import", str(malformed), success=False)
+
+    long_line = root / "long-line.jsonl"
+    long_line.write_text(
+        '{"type":"tiny-convmem-manifest","format_version":2}\n' + "x" * (1024 * 1024 + 4097),
+        encoding="utf-8",
+    )
+    run(server, root / "long-data", "--import", str(long_line), success=False)
+
+    atomic_target = root / "atomic.jsonl"
+    atomic_target.write_text("old backup", encoding="utf-8")
+    run(server, first, "--export", str(atomic_target), success=True)
+    assert "tiny-convmem-manifest" in atomic_target.read_text(encoding="utf-8")
+    assert not Path(str(atomic_target) + ".tmp").exists()
+
+    protected = root / "protected.jsonl"
+    protected.write_text("existing backup", encoding="utf-8")
+    Path(str(protected) + ".tmp").mkdir()
+    run(server, first, "--export", str(protected), success=False)
+    assert protected.read_text(encoding="utf-8") == "existing backup"
+
+
+if __name__ == "__main__":
+    main()