xinmi/scripts/test_weknora_category_etl.py

"""Category import regressions: tenant isolation, splits and tree ancestry."""
import copy
import unittest

from weknora_category_etl import TENANT, build_manifest, document_stem


def fixture():
    source = {
        "source_mutated": False, "relation_validation": {"valid": True}, "source": {"sha256": "fixture"},
        "mapping": {
            "knowledge_bases": [{"legacy_id": 1, "native_id": "plan-kb", "name": "Library"}],
            "categories": [
                {"legacy_id": 1, "native_id": "plan-root", "knowledge_base_id": "plan-kb", "parent_id": None, "name": "Root", "sort_order": 0},
                {"legacy_id": 2, "native_id": "plan-child", "knowledge_base_id": "plan-kb", "parent_id": "plan-root", "name": "Child", "sort_order": 1},
            ],
            "files": [{"legacy_id": 1, "knowledge_base_id": "plan-kb", "file_name": "Root__Child.jsonl", "category_id": "plan-child", "content_hash": "source-hash", "revision": 1}],
        },
    }
    state = {"tenant_id": TENANT, "knowledge_bases": {"1": {"id": "real-kb", "legacy_name": "Library", "complete": True, "failed": [], "uploaded": ["Root__Child.jsonl"]}}}
    target = {
        "tenant": {"id": TENANT, "status": "active"},
        "knowledge_bases": [{"id": "real-kb", "tenant_id": TENANT, "name": "Library"}],
        "tags": [], "relations": [],
        "files": [{"id": "doc", "tenant_id": TENANT, "knowledge_base_id": "real-kb", "file_name": "Root__Child.md", "file_hash": "target-hash", "file_size": 123, "parse_status": "completed", "folder_path": "", "custom_metadata": {"existing": "preserved"}}],
    }
    return source, state, target


class CategoryMigrationTests(unittest.TestCase):
    def test_parent_tag_and_folder_tree_use_actual_kb(self):
        result = build_manifest(*fixture())
        self.assertEqual(result["files"][0]["folder_path"], "Root/Child")
        self.assertEqual(result["files"][0]["metadata"]["existing"], "preserved")
        self.assertEqual(len(result["relations"]), 2)
        self.assertTrue(all(tag["knowledge_base_id"] == "real-kb" for tag in result["tags"]))

    def test_split_source_maps_both_parts(self):
        source, state, target = fixture()
        state["knowledge_bases"]["1"]["uploaded"] = ["Root__Child.part001.md", "Root__Child.part002.md"]
        file = target["files"][0]
        target["files"] = [{**file, "id": f"doc-{part}", "file_name": f"Root__Child.part{part:03}.md"} for part in (1, 2)]
        result = build_manifest(source, state, target)
        self.assertEqual(result["summary"]["source_files"], 1)
        self.assertEqual(result["summary"]["target_files"], 2)
        self.assertEqual(len(result["relations"]), 4)

    def test_partial_split_rejected(self):
        source, state, target = fixture()
        state["knowledge_bases"]["1"]["uploaded"] = ["Root__Child.part001.md", "Root__Child.part003.md"]
        file = target["files"][0]
        target["files"] = [{**file, "id": f"doc-{part}", "file_name": f"Root__Child.part{part:03}.md"} for part in (1, 3)]
        with self.assertRaisesRegex(ValueError, "split"):
            build_manifest(source, state, target)

    def test_missing_split_document_rejected(self):
        source, state, target = fixture()
        state["knowledge_bases"]["1"]["uploaded"] = ["Root__Child.part001.md", "Root__Child.part002.md"]
        target["files"][0]["file_name"] = "Root__Child.part001.md"
        with self.assertRaisesRegex(ValueError, "missing"):
            build_manifest(source, state, target)

    def test_cross_tenant_and_wrong_mapping_rejected(self):
        for item in ("tenant", "knowledge_base", "file"):
            with self.subTest(item=item):
                source, state, target = fixture()
                if item == "tenant":
                    target["tenant"]["id"] += 1
                elif item == "knowledge_base":
                    target["knowledge_bases"][0]["tenant_id"] += 1
                else:
                    target["files"][0]["tenant_id"] += 1
                with self.assertRaises(ValueError):
                    build_manifest(source, state, target)

    def test_same_name_does_not_adopt_existing_tag(self):
        source, state, target = fixture()
        target["tags"] = [{"id": "someone-elses-tag", "tenant_id": TENANT, "knowledge_base_id": "real-kb", "name": "Root", "sort_order": 0, "deleted_at": None}]
        with self.assertRaisesRegex(ValueError, "conflicts"):
            build_manifest(source, state, target)

    def test_second_plan_is_stable_and_preserves_user_metadata(self):
        source, state, target = fixture()
        first = build_manifest(source, state, target)
        target["tags"] = [{**t, "deleted_at": None} for t in first["tags"]]
        target["files"][0]["folder_path"] = first["files"][0]["folder_path"]
        target["files"][0]["custom_metadata"] = first["files"][0]["metadata"]
        second = build_manifest(source, state, target)
        self.assertEqual(first["tags"], second["tags"])
        self.assertEqual(first["relations"], second["relations"])
        self.assertEqual(first["files"][0]["metadata"], second["files"][0]["metadata"])

    def test_manual_changes_and_unmapped_files_are_not_overwritten(self):
        source, state, target = fixture()
        for change in ("folder", "unmapped", "duplicate"):
            changed = copy.deepcopy(target)
            if change == "folder":
                changed["files"][0]["folder_path"] = "User folder"
            elif change == "unmapped":
                changed["files"].append({**changed["files"][0], "id": "extra", "file_name": "Other.md"})
            else:
                changed["files"].append({**changed["files"][0], "id": "duplicate"})
            with self.subTest(change=change), self.assertRaises(ValueError):
                build_manifest(source, state, changed)

    def test_path_traversal_rejected(self):
        with self.assertRaises(ValueError):
            document_stem("../outside.md")
        source, state, target = fixture()
        source["mapping"]["categories"][0]["name"] = ".."
        with self.assertRaises(ValueError):
            build_manifest(source, state, target)


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