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", )