kg-agent/scripts/test_dms_production.py

"""Unit tests for template-driven DMS lifecycle extraction."""

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path
from urllib.parse import parse_qs, urlparse
from urllib.request import Request

from openpyxl import Workbook, load_workbook

from step0_pre_prepare import DmsAuthSettings
from step1_data_aggregation.dms_client import DmsClient, DmsReadError
from step1_data_aggregation.dms_pipeline import DATA_SHEET, _excel_value, build_production_data
from step1_data_aggregation.template_mapping import (
    DMS_FIELD,
    DMS_MODEL,
    MAPPING_SHEET,
    TEMPLATE_FIELD,
    load_template,
)


class StaticTokenManager:
    def call_with_refresh(self, operation):
        return operation("test-token")


class FakeDmsTransport:
    def __init__(self) -> None:
        self.content_requests: list[dict[str, list[str]]] = []

    def __call__(self, request: Request, timeout: float) -> bytes:
        path = urlparse(request.full_url).path
        if path == "/static/config/config.js":
            return b'window.config={dmsDataProxy:"/proxy_dms"}'
        if path == "/proxy_dms/column/getColumnList":
            return self.response(
                {
                    "code": 200,
                    "content": [
                        {
                            "title": "Group",
                            "columnList": [
                                {
                                    "id": 11,
                                    "title": "Model",
                                    "modelId": 22,
                                    "modelName": "model_table",
                                    "tag": "model_tag",
                                    "type": 1,
                                }
                            ],
                        },
                        {
                            "title": "Other",
                            "columnList": [
                                {
                                    "id": 33,
                                    "title": "Model",
                                    "modelId": 44,
                                    "modelName": "other_model",
                                    "type": 1,
                                }
                            ],
                        },
                    ],
                }
            )
        if path == "/proxy_dms/model/getModelById":
            model_id = parse_qs(urlparse(request.full_url).query)["modelId"]
            if model_id != ["22"]:
                return self.response({"code": 500, "message": "unexpected model"})
            return self.response(
                {
                    "code": 200,
                    "content": {
                        "fieldList": json.dumps(
                            {
                                "c_code": {"name": "c_code"},
                                "c_name": {"name": "c_name"},
                            }
                        )
                    },
                }
            )
        if path == "/proxy_dms/content/selectContentList":
            form = parse_qs((request.data or b"").decode("utf-8"))
            self.content_requests.append(form)
            page = int(form["page"][0])
            state = int(form["states"][0])
            if state == 0:
                data = (
                    [
                        {"c_code": "P001", "c_name": "Alpha"},
                        {"c_code": "P002", "c_name": "=not-a-formula"},
                    ]
                    if page == 0
                    else [{"c_code": "P003", "c_name": "Gamma"}]
                )
                count = 3
            elif state == 1:
                return self.response({"code": 202, "message": "data does not exist"})
            else:
                data = [{"c_code": f"P00{state + 3}", "c_name": f"State-{state}"}]
                count = 1
            return self.response({"code": 200, "content": {"count": count, "data": data}})
        raise AssertionError(f"unexpected request: {request.full_url}")

    @staticmethod
    def response(payload: dict[str, object]) -> bytes:
        return json.dumps(payload).encode("utf-8")


def make_template(path: Path) -> None:
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = MAPPING_SHEET
    sheet.append([TEMPLATE_FIELD, DMS_MODEL, DMS_FIELD])
    sheet.append(["Code", "Group-Model", "c_code"])
    sheet.append(["Name", "Group-Model", "c_name"])
    workbook.save(path)
    workbook.close()


class TemplateMappingTests(unittest.TestCase):
    def test_loads_mapping_rows_in_order(self) -> None:
        with tempfile.TemporaryDirectory() as raw_dir:
            path = Path(raw_dir) / "project.xlsx"
            make_template(path)
            template = load_template(path)
            self.assertEqual(template.name, "project")
            self.assertEqual(template.model_references, ("Group-Model",))
            self.assertEqual(
                [field.template_field for field in template.fields],
                ["Code", "Name"],
            )


class DmsClientTests(unittest.TestCase):
    def settings(self) -> DmsAuthSettings:
        return DmsAuthSettings(
            dms_base_url="https://dms.example/dms",
            oauth_api_url="https://oauth.example/oauth",
        )

    def test_timestamp_is_exported_as_shanghai_time_text(self) -> None:
        self.assertEqual(
            _excel_value(1769616000000, "timestamp"),
            "2026-01-29 00:00:00",
        )

    def test_discovers_proxy_resolves_exact_path_and_reads_all_states(self) -> None:
        transport = FakeDmsTransport()
        client = DmsClient(
            self.settings(),
            token_manager=StaticTokenManager(),
            transport=transport,
        )
        column = client.resolve_columns(["Group-Model"])["Group-Model"]
        self.assertEqual((column.column_id, column.model_id), (11, 22))
        self.assertEqual(client.get_model_fields(22), {"c_code", "c_name"})
        rows = client.fetch_rows(column, page_size=2)
        self.assertEqual([row["c_code"] for row in rows], ["P001", "P002", "P003", "P005", "P006"])
        self.assertEqual(len(transport.content_requests), 5)
        self.assertTrue(
            [request["states"][0] for request in transport.content_requests] == ["0", "0", "1", "2", "3"]
        )

    def test_rejects_ambiguous_leaf_title(self) -> None:
        client = DmsClient(
            self.settings(),
            token_manager=StaticTokenManager(),
            transport=FakeDmsTransport(),
        )
        with self.assertRaisesRegex(DmsReadError, "ambiguous"):
            client.resolve_columns(["Model"])


class ProductionPipelineTests(unittest.TestCase):
    def test_builds_valid_workbook_and_report(self) -> None:
        with tempfile.TemporaryDirectory() as raw_dir:
            root = Path(raw_dir)
            template_dir = root / "templates"
            output_dir = root / "production"
            template_dir.mkdir()
            make_template(template_dir / "project.xlsx")
            transport = FakeDmsTransport()
            client = DmsClient(
                DmsClientTests().settings(),
                token_manager=StaticTokenManager(),
                transport=transport,
            )
            results = build_production_data(
                client,
                template_dir=template_dir,
                output_dir=output_dir,
                page_size=2,
            )
            self.assertEqual(results[0].content_rows, 5)
            self.assertEqual(results[0].state_rows, {0: 3, 1: 0, 2: 1, 3: 1})
            output_path = output_dir / "project.xlsx"
            workbook = load_workbook(output_path, data_only=False)
            try:
                self.assertEqual(workbook.sheetnames, [DATA_SHEET])
                sheet = workbook[DATA_SHEET]
                self.assertEqual(
                    [cell.value for cell in sheet[1]],
                    ["Code", "Name"],
                )
                self.assertEqual(sheet.max_row, 6)
                self.assertEqual(sheet["B3"].value, "=not-a-formula")
                self.assertEqual(sheet["B3"].data_type, "s")
            finally:
                workbook.close()
            report = json.loads(
                (output_dir / "fetch_report.json").read_text(encoding="utf-8")
            )
            self.assertEqual(report["content_states"], [0, 1, 2, 3])
            self.assertEqual(report["templates"][0]["state_rows"], {"0": 3, "1": 0, "2": 1, "3": 1})
            self.assertEqual(report["templates"][0]["content_rows"], 5)


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