from __future__ import annotations import json from collections.abc import Iterator from contextlib import contextmanager from decimal import Decimal from typing import Any import pytest from app.modules.data_process.store import ( ConflictError, DataProcessStore, DataProcessStoreError, InvalidStateError, _decode_row, _preview_config_changed, _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]], datasets: list[dict[str, Any]] | None = None, ): self.results = results self.datasets: list[dict[str, Any]] = datasets or [] 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 params is not None: placeholder_count = normalized.count("%s") assert placeholder_count == len(params), ( f"SQL placeholder count {placeholder_count} does not match " f"parameter count {len(params)}" ) if normalized.startswith("SELECT * FROM data_process_results"): return _Result(rows=self.results) if normalized.startswith("SELECT * FROM datasets WHERE source='task'"): assert "source_task_id=%s" in normalized assert "source_task_id IS NULL AND task_id=%s" in normalized assert "deleted_at IS NULL" in normalized source_task_id, legacy_task_id = params return _Result( rows=[ item for item in self.datasets if item.get("source") == "task" and item.get("deleted_at") is None and ( item.get("source_task_id") == source_task_id or ( item.get("source_task_id") is None and item.get("task_id") == legacy_task_id ) ) ] ) if normalized.startswith("INSERT INTO datasets"): dataset = { "id": params[0], "name": params[1], "type": params[2], "source": "task", "task_id": params[4], "source_task_id": params[5], "count": params[8], "record_count": params[9], "metadata": params[11], "deleted_at": None, } 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 [] class _RegenerationConnection: def __init__( self, task: dict[str, Any], datasets: list[dict[str, Any]] | None = None, ) -> None: self.task = task self.datasets = datasets or [ { "id": "dataset_train", "source": "task", "task_id": task["id"], "source_task_id": task["id"], "deleted_at": None, }, { "id": "dataset_validation", "source": "task", "task_id": task["id"], "source_task_id": task["id"], "deleted_at": None, }, { "id": "dataset_test", "source": "task", "task_id": task["id"], "source_task_id": task["id"], "deleted_at": None, }, ] self.sources = [{"id": "source_1"}] self.previews = [{"id": "preview_1"}] self.results = [{"id": "result_1"}] def execute(self, sql: str, params: Any = None) -> _Result: normalized = " ".join(sql.split()) if params is not None: assert normalized.count("%s") == len(params) if normalized.startswith("UPDATE datasets SET source_task_id="): source_task_id, _, legacy_task_id = params for dataset in self.datasets: if ( dataset.get("source") == "task" and dataset.get("source_task_id") is None and dataset.get("task_id") == legacy_task_id and dataset.get("deleted_at") is None ): dataset["source_task_id"] = source_task_id return _Result() if normalized.startswith("SELECT EXISTS("): assert "source_task_id=%s" in normalized assert "source_task_id IS NULL AND task_id=%s" in normalized assert "deleted_at IS NULL" in normalized source_task_id, legacy_task_id = params exists = any( dataset.get("source") == "task" and dataset.get("deleted_at") is None and ( dataset.get("source_task_id") == source_task_id or ( dataset.get("source_task_id") is None and dataset.get("task_id") == legacy_task_id ) ) for dataset in self.datasets ) return _Result(row={"exists": exists}) if normalized.startswith("DELETE FROM data_process_results"): self.results.clear() return _Result() if normalized.startswith("DELETE FROM data_process_preview_items"): self.previews.clear() return _Result() if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_preview_items"): return _Result(row={"count": len(self.previews)}) if normalized.startswith("UPDATE data_process_tasks SET name="): self.task.update( { "name": params[0], "description": params[1], "config": params[2], "status": "pending", "progress": params[3], "output_dataset_id": None, "output_count": 0, "filtered_count": 0, "duplicate_count": 0, "error_count": 0, "failure_reason": None, "generation_run_id": None, "started_at": None, "completed_at": None, "updated_at": params[4], } ) return _Result(row=dict(self.task)) raise AssertionError(f"unexpected SQL: {normalized}") class _RegenerationStore(DataProcessStore): def __init__(self, conn: _RegenerationConnection): self._conn = conn @contextmanager def connect(self) -> Iterator[_RegenerationConnection]: yield self._conn def _task_in_connection( self, conn: Any, task_id: str, *, for_update: bool = False ) -> dict[str, Any]: assert for_update is True assert task_id == self._conn.task["id"] return dict(self._conn.task) class _TaskDetailConnection: def __init__( self, task: dict[str, Any], datasets: list[dict[str, Any]], ) -> None: self.task = task self.datasets = datasets def execute(self, sql: str, params: Any = None) -> _Result: normalized = " ".join(sql.split()) assert normalized.startswith("SELECT task.*") assert "dataset.source_task_id=task.id" in normalized assert "dataset.source_task_id IS NULL AND dataset.task_id=task.id" in normalized assert "dataset.deleted_at IS NULL" in normalized assert "dataset.metadata::jsonb" not in normalized assert "WHEN 'train' THEN 'train'" in normalized assert "WHEN 'val' THEN 'validation'" in normalized assert "WHEN 'test' THEN 'test'" in normalized task_id = params[0] visible = [ dataset for dataset in self.datasets if dataset.get("source") == "task" and dataset.get("deleted_at") is None and ( dataset.get("source_task_id") == task_id or ( dataset.get("source_task_id") is None and dataset.get("task_id") == task_id ) ) ] return _Result( row={ **self.task, "output_datasets": json.dumps( [ { "id": item["id"], "name": item["name"], "type": item["type"], "count": item["count"], "dataset_split": item["dataset_split"], } for item in visible ] ), } ) class _TaskDetailStore(DataProcessStore): def __init__(self, conn: _TaskDetailConnection) -> None: self._conn = conn @contextmanager def connect(self) -> Iterator[_TaskDetailConnection]: yield self._conn class _TaskListConnection: def __init__(self) -> None: self.task = { **_regeneration_task( status="pending", output_dataset_id=None, output_count=17, ), "created_at": "2026-07-25T18:00:00Z", "deleted_at": None, } self.sources = [ {"task_id": "task-1", "deleted_at": None}, {"task_id": "task-1", "deleted_at": None}, {"task_id": "task-1", "deleted_at": "2026-07-26T00:00:00Z"}, ] def execute(self, sql: str, params: Any = None) -> _Result: normalized = " ".join(sql.split()) if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_tasks task"): assert "task.deleted_at IS NULL" in normalized return _Result(row={"count": 1}) assert normalized.startswith("SELECT task.*") assert "source_file.task_id=task.id" in normalized assert "source_file.deleted_at IS NULL" in normalized source_file_count = sum( item["task_id"] == self.task["id"] and item["deleted_at"] is None for item in self.sources ) return _Result( rows=[ { **self.task, "source_file_count": source_file_count, } ] ) class _TaskListStore(DataProcessStore): def __init__(self, conn: _TaskListConnection) -> None: self._conn = conn @contextmanager def connect(self) -> Iterator[_TaskListConnection]: yield self._conn class _StartGenerationConnection: def __init__(self) -> None: self.task = { **_regeneration_task( status="pending", output_dataset_id=None, output_count=28, ), "generation_run_id": None, } self.results = [{"id": "old-result"}] def execute(self, sql: str, params: Any = None) -> _Result: normalized = " ".join(sql.split()) if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_preview_items"): return _Result(row={"count": 1}) if normalized.startswith("DELETE FROM data_process_results"): self.results.clear() return _Result() assert normalized.startswith("UPDATE data_process_tasks SET status='running'") assert "output_count=0" in normalized self.task.update( { "status": "running", "progress": 30, "output_count": 0, "failure_reason": None, "started_at": params[0], "completed_at": None, "filtered_count": 0, "duplicate_count": 0, "error_count": 0, "generation_run_id": params[1], "updated_at": params[2], } ) return _Result(row=dict(self.task)) class _StartGenerationStore(DataProcessStore): def __init__(self, conn: _StartGenerationConnection) -> None: self._conn = conn @contextmanager def connect(self) -> Iterator[_StartGenerationConnection]: yield self._conn def _task_in_connection( self, conn: Any, task_id: str, *, for_update: bool = False ) -> dict[str, Any]: assert task_id == "task-1" assert for_update is True return dict(self._conn.task) 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_decode_row_decodes_aggregated_output_datasets_json() -> None: decoded = _decode_row( { "id": "task-1", "output_datasets": '[{"id":"dataset_train","type":"train"}]', } ) assert decoded == { "id": "task-1", "output_datasets": [{"id": "dataset_train", "type": "train"}], } @pytest.mark.parametrize( ("process_type", "current", "next_config", "expected"), [ ( "structured", {"preprocess_options": ["deduplicate"]}, {"preprocess_options": ["deduplicate"], "temperature": 0.2}, False, ), ( "structured", {"preprocess_options": ["deduplicate"]}, {"preprocess_options": ["clean_invalid"]}, True, ), ( "structured", {"preprocess_options": ["a", "b"]}, {"preprocessOptions": ["b", "a", "a"], "chunk_size": 2048}, False, ), ( "unstructured", {"chunk_method": "fixed"}, {"chunk_method": "fixed", "generation_prompt": "new"}, False, ), ( "unstructured", {"chunk_method": "fixed"}, {"chunk_method": "semantic"}, True, ), ( "unstructured", {"chunk_size": 800, "chunk_overlap": 100}, {"chunk_size": 900, "chunk_overlap": 100}, True, ), ( "unstructured", {"min_chunk_size": 100}, {"min_chunk_size": 120}, True, ), ( "unstructured", {"semantic_breakpoint_percentile": 95}, {"semantic_breakpoint_percentile": 90}, True, ), ( "unstructured", {}, { "preserve_tables": True, "preserve_code_blocks": True, "preserve_lists": True, }, False, ), ( "unstructured", {"preserve_tables": False}, {"preserve_tables": True}, True, ), ( "unstructured", {"preserve_code_blocks": False}, {"preserve_code_blocks": True}, True, ), ( "unstructured", {"preserve_lists": False}, {"preserve_lists": True}, True, ), ( "unstructured", {"preprocess_options": ["deduplicate"]}, {"preprocess_options": ["clean_invalid"]}, True, ), ], ) def test_regeneration_preview_invalidation_matrix( process_type: str, current: dict[str, Any], next_config: dict[str, Any], expected: bool, ) -> None: assert _preview_config_changed(process_type, current, next_config) is expected def _regeneration_task(**updates: Any) -> dict[str, Any]: task = { "id": "task-1", "name": "原任务", "description": "", "process_type": "unstructured", "config": {"chunk_method": "fixed", "temperature": 0.7}, "status": "completed", "progress": 100, "output_dataset_id": "dataset_train", "output_count": 28, "filtered_count": 1, "duplicate_count": 1, "error_count": 0, "failure_reason": None, "generation_run_id": None, "started_at": "2026-07-25T18:00:00Z", "completed_at": "2026-07-25T18:05:00Z", "updated_at": "2026-07-25T18:05:00Z", } task.update(updates) return task def test_list_tasks_exposes_source_and_generation_counts() -> None: page = _TaskListStore(_TaskListConnection()).list_tasks(page=1, page_size=20) assert page["total"] == 1 assert page["page"] == 1 assert page["page_size"] == 20 item = page["items"][0] assert item["status"] == "pending" assert item["source_file_count"] == 2 assert item["output_count"] == 17 assert item["output_dataset_id"] is None def test_start_generation_clears_previous_output_count() -> None: conn = _StartGenerationConnection() task = _StartGenerationStore(conn).start_generation("task-1") assert task["status"] == "running" assert task["output_count"] == 0 assert conn.results == [] def _legacy_published_datasets(task_id: str = "task-1") -> list[dict[str, Any]]: specs = ( ("dataset_train", "制度问答-训练集", "train", "train", 22), ("dataset_validation", "制度问答-验证集", "val", "validation", 3), ("dataset_test", "制度问答-测试集", "test", "test", 3), ) dataset_ids = {split: dataset_id for dataset_id, _, _, split, _ in specs} return [ { "id": dataset_id, "name": name, "type": dataset_type, "source": "task", "task_id": task_id, "source_task_id": None, "count": count, "record_count": count, "dataset_split": split, "metadata": json.dumps( { "base_dataset_name": "制度问答", "dataset_split": split, "split_dataset_ids": dataset_ids, } ), "deleted_at": None, } for dataset_id, name, dataset_type, split, count in specs ] def test_prepare_regeneration_preserves_outputs_sources_and_generation_only_preview() -> None: conn = _RegenerationConnection(_regeneration_task()) original_datasets = list(conn.datasets) original_sources = list(conn.sources) result = _RegenerationStore(conn).prepare_regeneration( "task-1", { "name": "新任务名", "description": "更换生成参数", "process_type": "unstructured", "config": {"chunk_method": "fixed", "temperature": 0.2}, "expected_updated_at": "2026-07-25T18:05:00Z", }, ) assert result["preview_invalidated"] is False assert result["published_outputs_preserved"] is True assert result["task"]["output_dataset_id"] is None assert result["task"]["status"] == "pending" assert result["task"]["progress"] == 20 assert result["task"]["output_count"] == 0 assert result["task"]["started_at"] is None assert result["task"]["completed_at"] is None assert conn.results == [] assert conn.previews == [{"id": "preview_1"}] assert conn.datasets == original_datasets assert conn.sources == original_sources def test_prepare_regeneration_backfills_and_keeps_legacy_task_datasets_visible() -> None: legacy_datasets = _legacy_published_datasets() # 历史 metadata 可能不是合法 JSON，详情查询不能再依赖 metadata::jsonb。 legacy_datasets[0]["metadata"] = "{legacy-invalid-json" deleted_dataset = { **legacy_datasets[0], "id": "dataset_deleted", "name": "已删除训练集", "deleted_at": "2026-07-25T20:00:00Z", } unrelated_dataset = { **legacy_datasets[0], "id": "dataset_unrelated", "name": "其他任务训练集", "task_id": "task-other", } conn = _RegenerationConnection( _regeneration_task(), [*legacy_datasets, deleted_dataset, unrelated_dataset], ) before = _TaskDetailStore(_TaskDetailConnection(conn.task, conn.datasets)).get_task( "task-1" ) assert [item["id"] for item in before["output_datasets"]] == [ "dataset_train", "dataset_validation", "dataset_test", ] result = _RegenerationStore(conn).prepare_regeneration( "task-1", { "name": "原任务", "description": "", "process_type": "unstructured", "config": {"chunk_method": "fixed", "temperature": 0.2}, "expected_updated_at": "2026-07-25T18:05:00Z", }, ) assert result["published_outputs_preserved"] is True assert result["task"]["output_dataset_id"] is None assert len(conn.datasets) == 5 assert all( item["source_task_id"] == "task-1" for item in conn.datasets[:3] ) assert deleted_dataset["source_task_id"] is None assert unrelated_dataset["source_task_id"] is None after = _TaskDetailStore(_TaskDetailConnection(conn.task, conn.datasets)).get_task( "task-1" ) assert [item["id"] for item in after["output_datasets"]] == [ "dataset_train", "dataset_validation", "dataset_test", ] def test_prepare_regeneration_deletes_preview_when_chunk_configuration_changes() -> None: conn = _RegenerationConnection(_regeneration_task(output_dataset_id=None)) result = _RegenerationStore(conn).prepare_regeneration( "task-1", { "name": "原任务", "description": "", "process_type": "unstructured", "config": {"chunk_method": "semantic", "temperature": 0.7}, "expected_updated_at": "2026-07-25T18:05:00Z", }, ) assert result["preview_invalidated"] is True assert result["published_outputs_preserved"] is True assert result["task"]["progress"] == 0 assert conn.previews == [] assert conn.results == [] assert len(conn.datasets) == 3 @pytest.mark.parametrize( ("task_updates", "payload_updates", "error_type", "message"), [ ( {"status": "running"}, {}, ConflictError, "running task cannot be prepared", ), ( {}, {"expected_updated_at": "2026-07-25T17:00:00Z"}, ConflictError, "modified by another request", ), ( {}, {"process_type": "structured"}, InvalidStateError, "process_type cannot be changed", ), ], ) def test_prepare_regeneration_rejects_running_stale_and_type_change_without_mutation( task_updates: dict[str, Any], payload_updates: dict[str, Any], error_type: type[Exception], message: str, ) -> None: conn = _RegenerationConnection(_regeneration_task(**task_updates)) payload = { "name": "原任务", "description": "", "process_type": "unstructured", "config": {"chunk_method": "fixed"}, "expected_updated_at": "2026-07-25T18:05:00Z", **payload_updates, } with pytest.raises(error_type, match=message): _RegenerationStore(conn).prepare_regeneration("task-1", payload) assert conn.results == [{"id": "result_1"}] assert conn.previews == [{"id": "preview_1"}] assert conn.task["output_dataset_id"] == "dataset_train" 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_reuses_legacy_task_id_only_split_datasets() -> 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(10) ] legacy_datasets = _legacy_published_datasets() original_ids = [item["id"] for item in legacy_datasets] conn = _PublishConnection(results, legacy_datasets) published = _PublishStore(conn).publish( "task-1", { "dataset_name": "不会创建新数据集", "storage_type": "local", "format": "alpaca_jsonl", "split": {"train": 80, "validation": 10, "test": 10}, }, ) assert published["created"] is False assert [item["id"] for item in published["datasets"]] == original_ids assert [item["id"] for item in conn.datasets] == original_ids 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", )