252 lines
8.1 KiB
Python
252 lines
8.1 KiB
Python
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",
|
|
)
|