fix(data-process): 发布三个独立切分数据集

backend/tests/test_data_process_store.py

@@ -1,16 +1,112 @@
 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(
         {
@@ -22,6 +118,93 @@ def test_decode_row_serializes_postgres_numeric_values_as_json_numbers() -> None
     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"