YG_FT/backend/tests/test_data_process_store.py

from __future__ import annotations

from contextlib import contextmanager
from decimal import Decimal
from typing import Any, Iterator

import pytest

from app.modules.data_process.store import (
    DataProcessStore,
    DataProcessStoreError,
    _decode_row,
    _source_storage_descriptor,
)


class _Result:
    def __init__(self, *, row: dict[str, Any] | None = None, rows: list[dict[str, Any]] | None = None):
        self.row = row
        self.rows = rows or []

    def fetchone(self) -> dict[str, Any] | None:
        return self.row

    def fetchall(self) -> list[dict[str, Any]]:
        return self.rows


class _PublishConnection:
    def __init__(self, results: list[dict[str, Any]]):
        self.results = results
        self.datasets: list[dict[str, Any]] = []
        self.files: list[dict[str, Any]] = []
        self.records: list[dict[str, Any]] = []

    def execute(self, sql: str, params: Any = None) -> _Result:
        normalized = " ".join(sql.split())
        if normalized.startswith("SELECT * FROM data_process_results"):
            return _Result(rows=self.results)
        if normalized.startswith("SELECT * FROM datasets WHERE source_task_id"):
            return _Result(rows=self.datasets)
        if normalized.startswith("INSERT INTO datasets"):
            dataset = {
                "id": params[0],
                "name": params[1],
                "type": params[2],
                "count": params[8],
                "record_count": params[9],
                "metadata": params[11],
            }
            self.datasets.append(dataset)
            return _Result(row=dataset)
        if normalized.startswith("UPDATE datasets SET name="):
            dataset = next(item for item in self.datasets if item["id"] == params[10])
            dataset.update(
                {
                    "name": params[0],
                    "type": params[1],
                    "count": params[5],
                    "record_count": params[6],
                    "metadata": params[8],
                }
            )
            return _Result(row=dataset)
        if normalized.startswith("DELETE FROM dataset_records WHERE dataset_id"):
            self.records = [item for item in self.records if item["dataset_id"] != params[0]]
        if normalized.startswith("DELETE FROM dataset_files WHERE dataset_id"):
            self.files = [item for item in self.files if item["dataset_id"] != params[0]]
        if normalized.startswith("INSERT INTO dataset_files"):
            self.files.append(
                {
                    "id": params[0],
                    "dataset_id": params[1],
                    "name": params[2],
                    "record_count": params[11],
                }
            )
        if normalized.startswith("INSERT INTO dataset_records"):
            self.records.append(
                {"dataset_id": params[1], "line_no": params[4], "split": params[5]}
            )
        return _Result()


class _PublishStore(DataProcessStore):
    def __init__(self, conn: _PublishConnection):
        self._conn = conn

    @contextmanager
    def connect(self) -> Iterator[_PublishConnection]:
        yield self._conn

    def _task_in_connection(self, conn: Any, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
        train_dataset = next(
            (item for item in self._conn.datasets if item["type"] == "train"), None
        )
        return {
            "id": task_id,
            "status": "completed",
            "description": "",
            "config": {},
            "output_dataset_id": train_dataset and train_dataset["id"],
        }

    @staticmethod
    def _source_ids(conn: Any, task_id: str) -> list[dict[str, Any]]:
        return []


def test_decode_row_serializes_postgres_numeric_values_as_json_numbers() -> None:
    decoded = _decode_row(
        {
            "progress": Decimal("100.00"),
            "duration_seconds": Decimal("389.000000"),
        }
    )

    assert decoded == {"progress": 100.0, "duration_seconds": 389.0}


def test_publish_creates_three_independent_datasets_with_exact_counts() -> None:
    results = [
        {
            "id": f"result-{index}",
            "status": "valid",
            "instruction": f"问题 {index}",
            "input": "",
            "output": f"答案 {index}",
            "preview_item_id": f"preview-{index}",
        }
        for index in range(28)
    ]
    conn = _PublishConnection(results)
    published = _PublishStore(conn).publish(
        "task-1",
        {
            "dataset_name": "制度问答",
            "storage_type": "local",
            "format": "alpaca_jsonl",
            "split": {"train": 80, "validation": 10, "test": 10},
        },
    )

    assert [(item["name"], item["type"], item["count"]) for item in conn.datasets] == [
        ("制度问答-训练集", "train", 22),
        ("制度问答-验证集", "val", 3),
        ("制度问答-测试集", "test", 3),
    ]
    assert len(conn.files) == 3
    assert {item["dataset_id"] for item in conn.files} == {
        item["id"] for item in conn.datasets
    }
    assert len(conn.records) == 28
    assert published["dataset"]["type"] == "train"
    assert len(published["datasets"]) == 3
    assert published["split_counts"] == {"train": 22, "validation": 3, "test": 3}

    original_ids = [item["id"] for item in conn.datasets]
    republished = _PublishStore(conn).publish(
        "task-1",
        {
            "dataset_name": "制度问答-训练集",
            "storage_type": "local",
            "format": "alpaca_jsonl",
            "split": {"train": 80, "validation": 10, "test": 10},
        },
    )
    assert [item["id"] for item in conn.datasets] == original_ids
    assert len(conn.datasets) == 3
    assert len(conn.files) == 3
    assert len(conn.records) == 28
    assert republished["created"] is False


def test_publish_keeps_all_three_datasets_when_a_small_split_is_empty() -> None:
    conn = _PublishConnection(
        [
            {
                "id": "result-only",
                "status": "valid",
                "instruction": "唯一问题",
                "input": "",
                "output": "唯一答案",
                "preview_item_id": "preview-only",
            }
        ]
    )

    published = _PublishStore(conn).publish(
        "task-small",
        {
            "dataset_name": "小样本",
            "storage_type": "local",
            "format": "alpaca_jsonl",
            "split": {"train": 80, "validation": 10, "test": 10},
        },
    )

    assert [(item["type"], item["count"]) for item in conn.datasets] == [
        ("train", 1),
        ("val", 0),
        ("test", 0),
    ]
    assert len(published["datasets"]) == 3
    assert len(conn.files) == 3


def test_source_storage_descriptor_accepts_owned_local_and_legacy_db_references() -> None:
    task_id = "dpt_task"
    source_file_id = "dpsf_source"
    local_reference = (
        f"local://data-process/{task_id}/{source_file_id}/v1/source%20100%25.csv"
    )

    reference, metadata = _source_storage_descriptor(
        {
            "storage_object_id": local_reference,
            "metadata": {"storage_backend": "spoofed", "content_type": "text/csv"},
        },
        task_id,
        source_file_id,
    )
    assert reference == local_reference
    assert metadata == {"storage_backend": "local", "content_type": "text/csv"}

    legacy_reference, legacy_metadata = _source_storage_descriptor(
        {"metadata": {"legacy": True}},
        task_id,
        source_file_id,
    )
    assert legacy_reference == f"db://data-process/{task_id}/{source_file_id}/v1"
    assert legacy_metadata == {"storage_backend": "database", "legacy": True}


@pytest.mark.parametrize(
    "reference",
    [
        "local://data-process/dpt_other/dpsf_source/v1/source.txt",
        "db://data-process/dpt_task/dpsf_other/v1",
        "/var/tmp/source.txt",
    ],
)
def test_source_storage_descriptor_rejects_unowned_or_unsupported_references(
    reference: str,
) -> None:
    with pytest.raises(DataProcessStoreError):
        _source_storage_descriptor(
            {"storage_object_id": reference},
            "dpt_task",
            "dpsf_source",
        )