YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(data-process): 完善后台生成与失败重试

Browse Source
This commit is contained in:
caoxiaozhu
2026-07-28 10:56:05 +08:00
parent 8a6a6574bb
commit 01d2e6c76a
178 changed files with 4472 additions and 595 deletions
Download Patch File Download Diff File Expand all files Collapse all files

794
backend/tests/test_data_process_api.py
View File

@@ -3,11 +3,13 @@ from __future__ import annotations
from copy import deepcopy
from io import BytesIO
from pathlib import Path
from threading import Barrier, Lock
from typing import Any
import psycopg
import pytest
from docx import Document as WordDocument
from fastapi import FastAPI
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from openpyxl import Workbook
@@ -31,6 +33,7 @@ class FakeDataProcessStore:
self.previews: dict[str, list[dict[str, Any]]] = {}
self.results: dict[str, list[dict[str, Any]]] = {}
self.datasets: dict[str, dict[str, Any]] = {}
self.models: dict[str, dict[str, Any]] = {}
self.regeneration_prepared: set[str] = set()
self.sequence = 0
@@ -73,6 +76,14 @@ class FakeDataProcessStore:
"error_count": 0,
"failure_reason": None,
"output_dataset_id": None,
"results_confirmed": False,
"workflow_step": "create",
"preview_status": "idle",
"preview_progress": 0,
"preview_run_id": None,
"preview_failure_reason": None,
"preview_total_files": 0,
"preview_completed_files": 0,
}
self.tasks[task_id] = task
self.sources[task_id] = []
@@ -105,6 +116,11 @@ class FakeDataProcessStore:
self.tasks[task_id].update(deepcopy(payload))
return self.get_task(task_id)
def update_workflow_step(self, task_id: str, workflow_step: str) -> dict[str, Any]:
self.get_task(task_id)
self.tasks[task_id]["workflow_step"] = workflow_step
return self.get_task(task_id)
def prepare_regeneration(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
task = self.tasks.get(task_id)
if task is None:
@@ -136,8 +152,6 @@ class FakeDataProcessStore:
def delete_task(self, task_id: str, **_: Any) -> None:
self.get_task(task_id)
if self.tasks[task_id]["status"] == "running":
raise InvalidStateError("running task must be stopped before deletion")
del self.tasks[task_id]
def list_source_files(self, task_id: str) -> list[dict[str, Any]]:
@@ -197,6 +211,13 @@ class FakeDataProcessStore:
"error_count": 0,
"failure_reason": None,
"generation_run_id": None,
"workflow_step": "upload",
"preview_status": "idle",
"preview_progress": 0,
"preview_run_id": None,
"preview_failure_reason": None,
"preview_total_files": 0,
"preview_completed_files": 0,
"started_at": None,
"completed_at": None,
}
@@ -254,6 +275,15 @@ class FakeDataProcessStore:
item for item in self.previews[task_id] if item["source_file_id"] != file_id
]
self.results[task_id] = []
self.tasks[task_id].update(
workflow_step="upload",
preview_status="idle",
preview_progress=0,
preview_run_id=None,
preview_failure_reason=None,
preview_total_files=0,
preview_completed_files=0,
)
def replace_preview_items(
self,
@@ -261,7 +291,10 @@ class FakeDataProcessStore:
items: list[dict[str, Any]],
*,
source_file_ids: list[str] | None = None,
preview_run_id: str | None = None,
) -> list[dict[str, Any]]:
if preview_run_id is not None and not self.preview_is_running(task_id, preview_run_id):
raise InvalidStateError("preview run is no longer active")
created = [
{"id": self._id("dpp"), "task_id": task_id, **deepcopy(item)} for item in items
]
@@ -276,8 +309,121 @@ class FakeDataProcessStore:
] + created
self.results[task_id] = []
self.tasks[task_id]["progress"] = 20
if preview_run_id is None:
file_count = len(source_file_ids or {item["source_file_id"] for item in created})
self.tasks[task_id].update(
workflow_step="preview",
preview_status="completed",
preview_progress=100,
preview_run_id=None,
preview_failure_reason=None,
preview_total_files=file_count,
preview_completed_files=file_count,
)
return deepcopy(created)
def start_preview(
self,
task_id: str,
*,
source_file_ids: list[str] | None = None,
) -> tuple[dict[str, Any], list[str]]:
task = self.tasks[task_id]
if task.get("preview_status") in {"queued", "running"}:
raise InvalidStateError("task cannot be edited while preview is running")
selected_ids = source_file_ids or [str(item["id"]) for item in self.sources[task_id]]
found = {str(item["id"]) for item in self.sources[task_id]}
missing = set(selected_ids) - found
if missing:
raise NotFoundError(f"source files not found: {', '.join(sorted(missing))}")
if not selected_ids:
raise InvalidStateError("at least one source file is required")
run_id = self._id("dpprun")
self.results[task_id] = []
task.update(
status="pending",
workflow_step="upload",
preview_status="queued",
preview_progress=0,
preview_run_id=run_id,
preview_failure_reason=None,
preview_total_files=len(selected_ids),
preview_completed_files=0,
results_confirmed=False,
)
return self.get_task(task_id), list(selected_ids)
def mark_preview_running(self, task_id: str, preview_run_id: str) -> bool:
task = self.tasks.get(task_id)
if not task or task.get("preview_status") != "queued" or task.get("preview_run_id") != preview_run_id:
return False
task["preview_status"] = "running"
return True
def preview_is_running(self, task_id: str, preview_run_id: str) -> bool:
task = self.tasks.get(task_id)
return bool(
task
and task.get("preview_status") in {"queued", "running"}
and task.get("preview_run_id") == preview_run_id
)
def update_preview_progress(
self,
task_id: str,
preview_run_id: str,
completed_files: int,
total_files: int,
) -> bool:
if not self.preview_is_running(task_id, preview_run_id):
return False
self.tasks[task_id]["preview_completed_files"] = completed_files
self.tasks[task_id]["preview_progress"] = completed_files / max(1, total_files) * 100
return True
def complete_preview(self, task_id: str, preview_run_id: str) -> bool:
if not self.preview_is_running(task_id, preview_run_id):
return False
task = self.tasks[task_id]
task.update(
workflow_step="preview",
preview_status="completed",
preview_progress=100,
preview_run_id=None,
preview_failure_reason=None,
preview_completed_files=task["preview_total_files"],
)
return True
def mark_preview_failed(
self,
task_id: str,
reason: str,
*,
preview_run_id: str,
) -> bool:
if not self.preview_is_running(task_id, preview_run_id):
return False
self.tasks[task_id].update(
preview_status="failed",
preview_run_id=None,
preview_failure_reason=reason,
)
return True
def preview_progress(self, task_id: str) -> dict[str, Any]:
task = self.get_task(task_id)
return {
"task_id": task_id,
"workflow_step": task["workflow_step"],
"preview_status": task["preview_status"],
"preview_progress": task["preview_progress"],
"preview_run_id": task["preview_run_id"],
"preview_failure_reason": task["preview_failure_reason"],
"preview_total_files": task["preview_total_files"],
"preview_completed_files": task["preview_completed_files"],
}
def list_preview_items(
self,
task_id: str,
@@ -348,15 +494,19 @@ class FakeDataProcessStore:
progress=30,
output_dataset_id=None,
output_count=0,
results_confirmed=False,
workflow_step="generate",
generation_run_id=self._id("dprun"),
)
self.regeneration_prepared.discard(task_id)
return self.get_task(task_id)
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
task = self.tasks.get(task_id)
return (
self.tasks[task_id]["status"] == "running"
and self.tasks[task_id].get("generation_run_id") == generation_run_id
bool(task)
and task["status"] == "running"
and task.get("generation_run_id") == generation_run_id
)
def update_generation_progress(
@@ -389,6 +539,7 @@ class FakeDataProcessStore:
status="completed",
progress=100,
output_count=len(results),
results_confirmed=False,
generation_run_id=None,
**counts,
)
@@ -416,6 +567,7 @@ class FakeDataProcessStore:
result = {key: task.get(key) for key in (
"status", "progress", "input_count", "output_count",
"filtered_count", "duplicate_count", "error_count", "failure_reason",
"results_confirmed",
)}
result["task_id"] = task["id"]
return result
@@ -488,6 +640,66 @@ class FakeDataProcessStore:
item["status"] = "valid"
return deepcopy(item)
def get_generation_model(self, model_id: str) -> dict[str, Any]:
model = self.models.get(model_id)
if not model:
raise NotFoundError("generation model not found")
return deepcopy(model)
def replace_generated_result(
self,
task_id: str,
result_id: str,
replacement: dict[str, Any],
*,
expected_updated_at: str,
) -> dict[str, Any]:
task = self.tasks[task_id]
if task["status"] != "completed" or task.get("workflow_step") != "results":
raise InvalidStateError("task is not editing generation results")
if task.get("results_confirmed") or task.get("output_dataset_id"):
raise InvalidStateError("confirmed or published results cannot be regenerated")
item = next((item for item in self.results[task_id] if item["id"] == result_id), None)
if not item:
raise NotFoundError("data process result not found")
if item["status"] != "invalid":
raise InvalidStateError("only an invalid result can be regenerated")
if expected_updated_at != item.get("updated_at"):
raise InvalidStateError("data process result was modified by another request")
for field in ("instruction", "input", "output"):
item[field] = replacement[field]
item[f"original_{field}"] = replacement[field]
item.update(
status=replacement["status"],
error=replacement.get("error"),
quality_score=deepcopy(replacement["quality_score"]),
updated_at="2026-07-27T22:00:00Z",
)
task["error_count"] = sum(
result["status"] == "invalid" for result in self.results[task_id]
)
return deepcopy(item)
def confirm_results(self, task_id: str) -> dict[str, Any]:
task = self.tasks[task_id]
if task["status"] != "completed":
raise InvalidStateError("only a completed task can confirm results")
if task.get("workflow_step") != "results":
raise InvalidStateError("workflow must be on results before confirmation")
invalid_count = sum(
1
for item in self.results[task_id]
if item["status"] == "invalid"
or not item["instruction"].strip()
or not item["output"].strip()
)
if invalid_count:
raise InvalidStateError(f"task contains {invalid_count} invalid results")
if not self.results[task_id]:
raise InvalidStateError("task has no results to confirm")
task["results_confirmed"] = True
return self.get_task(task_id)
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
task = self.tasks[task_id]
if task_id in self.regeneration_prepared:
@@ -508,6 +720,8 @@ class FakeDataProcessStore:
}
if task["status"] != "completed":
raise InvalidStateError("only a completed task can be published")
if not task.get("results_confirmed"):
raise InvalidStateError("results must be confirmed before publishing")
split_specs = (
("train", "训练集"),
("val", "验证集"),
@@ -625,6 +839,14 @@ def _minimal_pdf(text: str = "Hello PDF") -> bytes:
return _minimal_pdf_pages(text)
def test_outdated_data_process_schema_returns_actionable_503() -> None:
with pytest.raises(HTTPException) as captured, data_process_endpoint.api_errors():
raise psycopg.errors.UndefinedColumn("missing runtime column")
assert captured.value.status_code == 503
assert "schema is missing or out of date" in captured.value.detail["message"]
def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
created = client.post(
@@ -636,6 +858,7 @@ def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
},
)
assert created.status_code == 200
assert created.json()["data"]["results_confirmed"] is False
task_id = created.json()["data"]["id"]
source_content = (
@@ -687,8 +910,10 @@ def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
generated = client.post(f"/modelTF/data-process/{task_id}/generate")
assert generated.status_code == 200
assert generated.json()["data"]["results_confirmed"] is False
progress = client.get(f"/modelTF/data-process/{task_id}/progress")
assert progress.json()["data"]["status"] == "completed"
assert progress.json()["data"]["results_confirmed"] is False
result_page = client.get(f"/modelTF/data-process/{task_id}/results").json()["data"]
assert result_page["total"] == 2
keyword_page = client.get(
@@ -719,6 +944,17 @@ def test_data_process_full_contract_without_database(tmp_path: Path) -> None:
assert restored.json()["data"]["status"] == "valid"
assert store.tasks[task_id]["error_count"] == 0
workflow = client.put(
f"/modelTF/data-process/{task_id}/workflow-step",
json={"workflow_step": "results"},
)
assert workflow.status_code == 200
confirmed = client.post(
f"/modelTF/data-process/{task_id}/confirm-results"
)
assert confirmed.status_code == 200
assert confirmed.json()["data"]["results_confirmed"] is True
publish_payload = {"dataset_name": "客服问答清洗集"}
first_publish = client.post(
f"/modelTF/data-process/{task_id}/publish", json=publish_payload
@@ -761,6 +997,120 @@ def test_task_list_exposes_document_and_generation_counts(
assert item["output_dataset_id"] is None
def test_workflow_step_update_is_independent_and_validated(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "步骤持久化", "process_type": "structured", "config": {}},
).json()["data"]["id"]
store.previews[task_id] = [{"id": "preview-1"}]
store.results[task_id] = [{"id": "result-1"}]
store.tasks[task_id].update(status="running", generation_run_id="run-1")
updated = client.put(
f"/modelTF/data-process/{task_id}/workflow-step",
json={"workflow_step": "generate"},
)
assert updated.status_code == 200
assert updated.json()["data"]["workflow_step"] == "generate"
assert store.previews[task_id] == [{"id": "preview-1"}]
assert store.results[task_id] == [{"id": "result-1"}]
assert store.tasks[task_id]["status"] == "running"
assert client.put(
f"/modelTF/data-process/{task_id}/workflow-step",
json={"workflow_step": "unknown"},
).status_code == 422
def test_background_preview_persists_progress_and_items(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "后台切分", "process_type": "structured", "config": {}},
).json()["data"]["id"]
source = client.post(
f"/modelTF/data-process/{task_id}/source-files",
files={
"files": (
"one.jsonl",
b'{"question":"What is one?","answer":"One."}\n',
"application/jsonl",
)
},
).json()["data"]["files"][0]
started = client.post(
f"/modelTF/data-process/{task_id}/preview/start",
json={"source_file_ids": [source["id"]]},
)
assert started.status_code == 202
assert started.json()["data"]["preview_status"] == "queued"
assert started.json()["data"]["preview_run_id"]
progress = client.get(
f"/modelTF/data-process/{task_id}/preview/progress"
).json()["data"]
assert progress == {
"task_id": task_id,
"workflow_step": "preview",
"preview_status": "completed",
"preview_progress": 100.0,
"preview_run_id": None,
"preview_failure_reason": None,
"preview_total_files": 1,
"preview_completed_files": 1,
}
assert client.get(
f"/modelTF/data-process/{task_id}/preview"
).json()["data"]["total"] == 1
def test_stale_preview_worker_cannot_replace_new_run(tmp_path: Path) -> None:
_, store, storage = make_client(tmp_path)
task = store.create_task(
{"name": "切分代次", "process_type": "structured", "config": {}}
)
task_id = str(task["id"])
store.sources[task_id] = [{"id": "source-1"}]
first, source_ids = store.start_preview(task_id, source_file_ids=["source-1"])
first_run_id = str(first["preview_run_id"])
store.tasks[task_id].update(preview_status="cancelled", preview_run_id=None)
second, _ = store.start_preview(task_id, source_file_ids=["source-1"])
data_process_endpoint._run_preview(
store,
storage,
task_id,
first_run_id,
source_ids,
)
assert store.tasks[task_id]["preview_status"] == "queued"
assert store.tasks[task_id]["preview_run_id"] == second["preview_run_id"]
assert store.previews[task_id] == []
def test_delete_invalidates_active_generation_and_preview(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "删除运行任务", "process_type": "structured", "config": {}},
).json()["data"]["id"]
store.tasks[task_id].update(
status="running",
generation_run_id="generation-1",
preview_status="running",
preview_run_id="preview-1",
)
deleted = client.delete(f"/modelTF/data-process/{task_id}")
assert deleted.status_code == 200
assert store.generation_is_running(task_id, "generation-1") is False
assert store.preview_is_running(task_id, "preview-1") is False
def test_task_detail_uses_returned_source_files_as_document_count(tmp_path: Path) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
@@ -811,6 +1161,404 @@ def test_generation_start_response_clears_previous_output_count(
assert store.tasks[task_id]["output_count"] == 0
def test_results_must_be_generated_and_valid_before_confirmation(
tmp_path: Path,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={"name": "结果确认门禁", "process_type": "structured", "config": {}},
).json()["data"]["id"]
pending = client.post(f"/modelTF/data-process/{task_id}/confirm-results")
assert pending.status_code == 409
store.tasks[task_id].update(status="completed", progress=100, workflow_step="generate")
store.results[task_id] = [
{
"id": "result_valid",
"status": "valid",
"instruction": "问题",
"input": "",
"output": "答案",
}
]
wrong_step = client.post(f"/modelTF/data-process/{task_id}/confirm-results")
assert wrong_step.status_code == 409
store.tasks[task_id]["workflow_step"] = "results"
store.results[task_id] = [
{
"id": "result_invalid",
"status": "invalid",
"instruction": "问题",
"input": "",
"output": "",
}
]
invalid = client.post(f"/modelTF/data-process/{task_id}/confirm-results")
assert invalid.status_code == 409
store.results[task_id][0].update(status="valid", output="答案")
confirmed = client.post(f"/modelTF/data-process/{task_id}/confirm-results")
assert confirmed.status_code == 200
assert confirmed.json()["data"]["results_confirmed"] is True
assert client.post(
f"/modelTF/data-process/{task_id}/confirm-results"
).status_code == 200
def test_invalid_result_can_be_regenerated_in_place(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
"name": "单条重生成",
"process_type": "structured",
"config": {
"generation_model_id": "model-1",
"output_type": "standard",
"min_output_length": 5,
"generation_retries": 5,
"request_timeout_seconds": 120,
},
},
).json()["data"]["id"]
store.tasks[task_id].update(
status="completed",
progress=100,
workflow_step="results",
results_confirmed=False,
error_count=1,
)
store.models["model-1"] = {
"id": "model-1",
"name": "测试模型",
"online_model_name": "test-model",
"api_url": "https://model.example/v1",
"api_key": "secret",
}
store.previews[task_id] = [{
"id": "preview-1",
"status": "original",
"original_content": "申请编号字段用于标识报销申请。",
"edited_content": "申请编号字段用于标识报销申请。",
}]
store.results[task_id] = [{
"id": "result-1",
"preview_item_id": "preview-1",
"instruction": "模型生成失败，请人工补充",
"input": "申请编号字段用于标识报销申请。",
"output": "",
"original_instruction": "模型生成失败，请人工补充",
"original_input": "申请编号字段用于标识报销申请。",
"original_output": "",
"status": "invalid",
"error": "model response is not valid JSON",
"split": "train",
"quality_score": {},
"updated_at": "2026-07-27T21:00:00Z",
}]
captured: dict[str, Any] = {}
def fake_generate(preview_items: Any, **kwargs: Any) -> list[dict[str, Any]]:
captured["preview_items"] = list(preview_items)
captured["qa_pairs_per_item"] = kwargs["qa_pairs_per_item"]
captured["config"] = kwargs["config"]
return [{
"id": "temporary-result",
"preview_item_id": "preview-1",
"instruction": "申请编号字段有什么作用？",
"input": "申请编号字段用于标识报销申请。",
"output": "申请编号字段用于唯一标识一笔报销申请。",
"original_instruction": "申请编号字段有什么作用？",
"original_input": "申请编号字段用于标识报销申请。",
"original_output": "申请编号字段用于唯一标识一笔报销申请。",
"status": "valid",
"error": None,
"split": "train",
}]
monkeypatch.setattr(data_process_endpoint, "generate_model_records", fake_generate)
response = client.post(
f"/modelTF/data-process/{task_id}/results/result-1/regenerate",
json={"expected_updated_at": "2026-07-27T21:00:00Z"},
)
assert response.status_code == 200
regenerated = response.json()["data"]
assert regenerated["id"] == "result-1"
assert regenerated["status"] == "valid"
assert regenerated["instruction"] == regenerated["original_instruction"]
assert regenerated["output"] == regenerated["original_output"]
assert store.tasks[task_id]["error_count"] == 0
assert captured["qa_pairs_per_item"] == 1
assert [item["id"] for item in captured["preview_items"]] == ["preview-1"]
assert captured["config"]["generation_retries"] == 0
assert captured["config"]["request_timeout_seconds"] == 60
def test_failed_result_regeneration_keeps_the_original_error(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
"name": "失败结果保留",
"process_type": "structured",
"config": {
"generation_model_id": "model-1",
"output_type": "standard",
},
},
).json()["data"]["id"]
store.tasks[task_id].update(
status="completed",
progress=100,
workflow_step="results",
results_confirmed=False,
error_count=1,
)
store.models["model-1"] = {
"id": "model-1",
"online_model_name": "test-model",
"api_url": "https://model.example/v1",
"api_key": "secret",
}
store.previews[task_id] = [{
"id": "preview-1",
"status": "original",
"original_content": "原始内容",
"edited_content": "原始内容",
}]
original_result = {
"id": "result-1",
"preview_item_id": "preview-1",
"instruction": "模型生成失败，请人工补充",
"input": "原始内容",
"output": "",
"status": "invalid",
"error": "first failure",
"updated_at": "2026-07-27T21:30:00Z",
}
store.results[task_id] = [original_result.copy()]
monkeypatch.setattr(
data_process_endpoint,
"generate_model_records",
lambda *args, **kwargs: [{
"status": "invalid",
"error": "second failure",
}],
)
response = client.post(
f"/modelTF/data-process/{task_id}/results/result-1/regenerate",
json={"expected_updated_at": "2026-07-27T21:30:00Z"},
)
assert response.status_code == 409
assert store.results[task_id] == [original_result]
assert store.tasks[task_id]["error_count"] == 1
def test_failed_results_can_be_regenerated_in_parallel_with_partial_success(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
"name": "批量重生成",
"process_type": "structured",
"config": {
"generation_model_id": "model-1",
"output_type": "standard",
"min_output_length": 5,
},
},
).json()["data"]["id"]
store.tasks[task_id].update(
status="completed",
progress=100,
workflow_step="results",
results_confirmed=False,
error_count=2,
)
store.models["model-1"] = {
"id": "model-1",
"online_model_name": "test-model",
"api_url": "https://model.example/v1",
"api_key": "secret",
}
store.previews[task_id] = [
{
"id": "preview-1",
"status": "original",
"original_content": "申请编号用于唯一标识一笔报销申请。",
"edited_content": "申请编号用于唯一标识一笔报销申请。",
},
{
"id": "preview-2",
"status": "original",
"original_content": "联系电话用于联系申请人。",
"edited_content": "联系电话用于联系申请人。",
},
]
original_results = [
{
"id": "result-1",
"preview_item_id": "preview-1",
"instruction": "模型生成失败，请人工补充",
"input": "申请编号用于唯一标识一笔报销申请。",
"output": "",
"original_instruction": "模型生成失败，请人工补充",
"original_input": "申请编号用于唯一标识一笔报销申请。",
"original_output": "",
"status": "invalid",
"error": "first failure",
"split": "train",
"quality_score": {},
"updated_at": "2026-07-28T09:00:00Z",
},
{
"id": "result-2",
"preview_item_id": "preview-2",
"instruction": "模型生成失败，请人工补充",
"input": "联系电话用于联系申请人。",
"output": "",
"original_instruction": "模型生成失败，请人工补充",
"original_input": "联系电话用于联系申请人。",
"original_output": "",
"status": "invalid",
"error": "first failure",
"split": "train",
"quality_score": {},
"updated_at": "2026-07-28T09:00:01Z",
},
]
store.results[task_id] = deepcopy(original_results)
barrier = Barrier(2, timeout=2)
activity_lock = Lock()
active_calls = 0
max_active_calls = 0
model_clients: list[Any] = []
def fake_generate(preview_items: Any, **kwargs: Any) -> list[dict[str, Any]]:
nonlocal active_calls, max_active_calls
preview = next(iter(preview_items))
with activity_lock:
active_calls += 1
max_active_calls = max(max_active_calls, active_calls)
model_clients.append(kwargs.get("client"))
try:
barrier.wait()
if preview["id"] == "preview-2":
return [{"status": "invalid", "error": "second failure"}]
return [{
"id": "temporary-result",
"preview_item_id": preview["id"],
"instruction": "申请编号有什么作用？",
"input": preview["edited_content"],
"output": "申请编号用于唯一标识一笔报销申请。",
"status": "valid",
"error": None,
"split": "train",
}]
finally:
with activity_lock:
active_calls -= 1
monkeypatch.setattr(data_process_endpoint, "generate_model_records", fake_generate)
response = client.post(
f"/modelTF/data-process/{task_id}/results/regenerate-batch",
json={
"items": [
{
"result_id": "result-1",
"expected_updated_at": "2026-07-28T09:00:00Z",
},
{
"result_id": "result-2",
"expected_updated_at": "2026-07-28T09:00:01Z",
},
],
},
)
assert response.status_code == 200
data = response.json()["data"]
assert data["total"] == 2
assert data["succeeded"] == 1
assert data["failed"] == 1
assert data["remaining_invalid_count"] == 1
assert [item["id"] for item in data["items"]] == ["result-1"]
assert data["failures"][0]["result_id"] == "result-2"
assert store.results[task_id][0]["status"] == "valid"
assert store.results[task_id][0]["id"] == "result-1"
assert store.results[task_id][1] == original_results[1]
assert store.tasks[task_id]["error_count"] == 1
assert max_active_calls == 2
assert all(client is not None for client in model_clients)
assert len({id(client) for client in model_clients}) == 1
@pytest.mark.parametrize(
"locked_state",
[
{"results_confirmed": True},
{"output_dataset_id": "dataset-published"},
],
ids=["confirmed", "published"],
)
def test_batch_result_regeneration_rejects_locked_tasks_before_model_call(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
locked_state: dict[str, Any],
) -> None:
client, store, _ = make_client(tmp_path)
task_id = client.post(
"/modelTF/data-process",
json={
"name": "批量重生成门禁",
"process_type": "structured",
"config": {"generation_model_id": "model-1"},
},
).json()["data"]["id"]
store.tasks[task_id].update(
status="completed",
workflow_step="results",
results_confirmed=False,
)
store.tasks[task_id].update(locked_state)
model_calls = 0
def fake_generate(*args: Any, **kwargs: Any) -> list[dict[str, Any]]:
nonlocal model_calls
model_calls += 1
return []
monkeypatch.setattr(data_process_endpoint, "generate_model_records", fake_generate)
response = client.post(
f"/modelTF/data-process/{task_id}/results/regenerate-batch",
json={
"items": [{
"result_id": "result-1",
"expected_updated_at": "2026-07-28T09:00:00Z",
}],
},
)
assert response.status_code == 409
assert model_calls == 0
def test_preview_build_replaces_only_selected_files_and_reports_file_counts(
tmp_path: Path,
) -> None:
@@ -1026,6 +1774,7 @@ def test_published_split_datasets_remain_in_detail_after_regeneration(
"status": "completed",
"updated_at": "2026-07-27T09:00:00Z",
"output_count": 1,
"results_confirmed": True,
}
)
store.results[task_id] = [
@@ -1440,6 +2189,41 @@ def test_reasoning_output_requires_generation_model() -> None:
assert failed["failure_reason"] == "思维链输出必须配置可用的数据生成模型"
def test_generation_failure_is_written_to_structured_log(caplog: pytest.LogCaptureFixture) -> None:
store = FakeDataProcessStore()
task = store.create_task(
{
"name": "生成失败日志",
"process_type": "structured",
"config": {"output_type": "reasoning"},
}
)
task_id = task["id"]
store.replace_preview_items(
task_id,
[
{
"source_file_id": None,
"original_content": "需要推理的来源内容",
"edited_content": "需要推理的来源内容",
"status": "manual",
}
],
)
started = store.start_generation(task_id, replace_existing=True)
with caplog.at_level("INFO", logger=data_process_endpoint.__name__):
data_process_endpoint._run_generation(
store,
task_id,
started["generation_run_id"],
)
messages = [record.getMessage() for record in caplog.records]
assert any("generation worker started" in message for message in messages)
assert any("generation failed" in message for message in messages)
def test_result_status_cannot_be_forged_by_client(tmp_path: Path) -> None:
client, _, _ = make_client(tmp_path)
task_id = client.post(

444
backend/tests/test_data_process_generation.py
View File

@@ -88,6 +88,344 @@ def test_generate_model_records_uses_prompt_auth_and_stable_split() -> None:
assert progress_updates == [(1, 1)]
def test_minimax_m3_uses_split_reasoning_and_completion_token_budget() -> None:
requests: list[dict[str, object]] = []
def handler(request: httpx.Request) -> httpx.Response:
payload = json.loads(request.content)
requests.append(payload)
return httpx.Response(
200,
json={
"choices": [{
"finish_reason": "stop",
"message": {
"reasoning_content": "模型内部思考不应混入业务 JSON",
"content": json.dumps({
"items": [{
"instruction": "申请编号有什么作用？",
"reasoning": "来源说明它用于标识报销申请。",
"answer": "它用于唯一标识一笔报销申请。",
}],
}, ensure_ascii=False),
},
}],
"output_sensitive": False,
"base_resp": {"status_code": 0, "status_msg": ""},
},
)
records = generate_model_records(
[{"id": "preview-minimax", "edited_content": "申请编号用于标识报销申请。"}],
model={
"name": "MiniMax",
"online_model_name": "MiniMax-M3",
"api_url": "https://api.minimaxi.com/v1",
},
config={
"output_type": "reasoning",
"json_mode": True,
"max_tokens": 1024,
"generation_retries": 0,
},
task_id="task-minimax",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert records[0]["status"] == "valid"
assert len(requests) == 1
assert requests[0]["reasoning_split"] is True
assert requests[0]["max_completion_tokens"] >= 4096
assert "max_tokens" not in requests[0]
assert "response_format" not in requests[0]
def test_minimax_m3_keeps_larger_configured_completion_budget() -> None:
requests: list[dict[str, object]] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(json.loads(request.content))
return httpx.Response(
200,
json={
"choices": [{
"message": {
"content": json.dumps({
"items": [{
"instruction": "问题",
"output": "这是满足测试要求的完整答案。",
}],
}, ensure_ascii=False),
},
}],
},
)
generate_model_records(
[{"id": "preview-minimax-budget", "edited_content": "来源正文"}],
model={
"online_model_name": "MiniMax-M3",
"api_url": "https://api.minimax.io/v1",
},
config={"max_tokens": 8192, "generation_retries": 0},
task_id="task-minimax-budget",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert requests[0]["max_completion_tokens"] == 8192
def test_minimax_m3_name_on_custom_proxy_keeps_generic_openai_parameters() -> None:
requests: list[dict[str, object]] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(json.loads(request.content))
return httpx.Response(
200,
json={
"choices": [{
"message": {
"content": json.dumps({
"items": [{
"instruction": "问题",
"output": "这是代理服务返回的完整答案。",
}],
}, ensure_ascii=False),
},
}],
},
)
generate_model_records(
[{"id": "preview-minimax-proxy", "edited_content": "来源正文"}],
model={
"online_model_name": "MiniMax-M3",
"api_url": "https://model-proxy.example/v1",
},
config={"max_tokens": 1024, "json_mode": True, "generation_retries": 0},
task_id="task-minimax-proxy",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert requests[0]["max_tokens"] == 1024
assert requests[0]["response_format"] == {"type": "json_object"}
assert "reasoning_split" not in requests[0]
assert "max_completion_tokens" not in requests[0]
def test_generate_model_records_extracts_json_surrounded_by_model_explanation() -> None:
content = "模型结果如下：\n```json\n" + json.dumps(
{
"items": [{
"instruction": "字段有什么作用？",
"output": "该字段用于唯一标识记录。",
}],
},
ensure_ascii=False,
) + "\n```\n生成完毕。"
client = httpx.Client(
transport=httpx.MockTransport(
lambda _: httpx.Response(
200,
json={"choices": [{"message": {"content": content}}]},
)
)
)
records = generate_model_records(
[{"id": "preview-explanation", "edited_content": "字段用于唯一标识记录。"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 0},
task_id="task-explanation",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=client,
)
assert records[0]["status"] == "valid"
assert records[0]["output"] == "该字段用于唯一标识记录。"
def test_generate_model_records_reports_token_truncation_instead_of_json_error() -> None:
client = httpx.Client(
transport=httpx.MockTransport(
lambda _: httpx.Response(
200,
json={
"choices": [{
"finish_reason": "length",
"message": {"content": ""},
}],
"output_sensitive": False,
},
)
)
)
records = generate_model_records(
[{"id": "preview-truncated", "edited_content": "来源正文"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 0},
task_id="task-truncated",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=client,
)
assert records[0]["status"] == "invalid"
assert "Token" in records[0]["error"]
assert "截断" in records[0]["error"]
def test_token_truncation_is_not_retried_even_when_json_looks_complete() -> None:
request_count = 0
content = json.dumps({
"items": [{
"instruction": "问题",
"output": "表面完整但服务端已声明截断。",
}],
}, ensure_ascii=False)
def handler(_: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
return httpx.Response(
200,
json={
"choices": [{
"finish_reason": "length",
"message": {"content": content},
}],
},
)
records = generate_model_records(
[{"id": "preview-length", "edited_content": "来源正文"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 5},
task_id="task-length",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 1
assert records[0]["status"] == "invalid"
assert "finish_reason=length" in records[0]["error"]
def test_sensitive_model_response_is_not_retried_or_saved() -> None:
request_count = 0
def handler(_: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
return httpx.Response(
200,
json={
"choices": [{
"finish_reason": "stop",
"message": {"content": "{}"},
}],
"output_sensitive": True,
"base_resp": {"status_code": 1027, "status_msg": "output sensitive"},
},
)
records = generate_model_records(
[{"id": "preview-sensitive", "edited_content": "来源正文"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 5},
task_id="task-sensitive",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 1
assert records[0]["status"] == "invalid"
assert "安全拦截" in records[0]["error"]
assert "1027" in records[0]["error"]
def test_empty_model_content_can_retry_then_succeed() -> None:
request_count = 0
def handler(_: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
if request_count == 1:
return httpx.Response(
200,
json={"choices": [{"finish_reason": "stop", "message": {"content": ""}}]},
)
return httpx.Response(
200,
json={
"choices": [{
"finish_reason": "stop",
"message": {
"content": json.dumps({
"items": [{
"instruction": "问题",
"output": "第二次请求返回了完整答案。",
}],
}, ensure_ascii=False),
},
}],
},
)
records = generate_model_records(
[{"id": "preview-empty-retry", "edited_content": "来源正文"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 1},
task_id="task-empty-retry",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 2
assert records[0]["status"] == "valid"
def test_multiple_top_level_json_documents_are_rejected_as_ambiguous() -> None:
first = json.dumps({
"items": [{"instruction": "问题一", "output": "答案一"}],
}, ensure_ascii=False)
second = json.dumps({
"items": [{"instruction": "问题二", "output": "答案二"}],
}, ensure_ascii=False)
client = httpx.Client(
transport=httpx.MockTransport(
lambda _: httpx.Response(
200,
json={"choices": [{"message": {"content": f"{first}\n{second}"}}]},
)
)
)
records = generate_model_records(
[{"id": "preview-ambiguous", "edited_content": "来源正文"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 0},
task_id="task-ambiguous",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=client,
)
assert records[0]["status"] == "invalid"
assert "多个 JSON" in records[0]["error"]
def test_generate_model_records_builds_reasoning_output_with_think_tags() -> None:
def handler(request: httpx.Request) -> httpx.Response:
payload = json.loads(request.content)
@@ -414,6 +752,112 @@ def test_generate_model_records_retries_short_batch_then_marks_it_invalid() -> N
assert "expected 10, got 1" in records[0]["error"]
def test_generate_model_records_does_not_retry_non_retryable_http_errors() -> None:
request_count = 0
def handler(_: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
return httpx.Response(401, json={"error": {"message": "unauthorized"}})
records = generate_model_records(
[{"id": "preview-auth", "edited_content": "来源内容"}],
model={
"api_url": "https://model.example/v1",
"online_model_name": "test-model",
"api_key": "invalid",
},
config={"generation_retries": 5},
task_id="task-auth",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 1
assert records[0]["status"] == "invalid"
assert "401" in records[0]["error"]
@pytest.mark.parametrize("status_code", [408, 425, 429, 500])
def test_generate_model_records_retries_retryable_http_statuses(
status_code: int,
) -> None:
request_count = 0
def handler(_: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
if request_count == 1:
return httpx.Response(status_code)
return httpx.Response(
200,
json={
"choices": [{
"message": {
"content": json.dumps({
"items": [{
"instruction": "来源内容是什么？",
"output": "这是用于验证可重试错误的来源内容。",
}],
}, ensure_ascii=False),
},
}],
},
)
records = generate_model_records(
[{"id": "preview-retryable", "edited_content": "来源内容"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 1},
task_id="task-retryable",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 2
assert records[0]["status"] == "valid"
def test_generate_model_records_retries_transient_network_errors() -> None:
request_count = 0
def handler(request: httpx.Request) -> httpx.Response:
nonlocal request_count
request_count += 1
if request_count == 1:
raise httpx.ConnectError("temporary connection failure", request=request)
return httpx.Response(
200,
json={
"choices": [{
"message": {
"content": json.dumps({
"items": [{
"instruction": "网络恢复了吗？",
"output": "临时连接错误后，第二次模型请求已经成功。",
}],
}, ensure_ascii=False),
},
}],
},
)
records = generate_model_records(
[{"id": "preview-network", "edited_content": "网络重试来源"}],
model={"name": "model", "api_url": "https://model.example/v1"},
config={"generation_retries": 1},
task_id="task-network",
split={"train": 100, "validation": 0, "test": 0},
qa_pairs_per_item=1,
client=httpx.Client(transport=httpx.MockTransport(handler)),
)
assert request_count == 2
assert records[0]["status"] == "valid"
@pytest.mark.parametrize("qa_pairs_per_item", [0, 51])
def test_generate_model_records_rejects_out_of_range_count(
qa_pairs_per_item: int,

34
backend/tests/test_data_process_migration.py
View File

@@ -2,7 +2,7 @@ from __future__ import annotations
from pathlib import Path
from app.modules.data_process.schema_cli import _target_label
from app.modules.data_process.schema_cli import REQUIRED_TASK_COLUMNS, _target_label
def test_runtime_migration_fails_fast_on_incompatible_schema() -> None:
@@ -18,6 +18,24 @@ def test_runtime_migration_fails_fast_on_incompatible_schema() -> None:
assert "requires 001_platform_runtime.sql first" in sql
assert "supports only the current TEXT runtime schema" in sql
assert "generation_run_id" in sql
assert "results_confirmed BOOLEAN NOT NULL DEFAULT TRUE" in sql
assert "WHERE status <> 'completed' AND results_confirmed=TRUE" in sql
assert "ADD COLUMN IF NOT EXISTS workflow_step VARCHAR(20)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_status VARCHAR(20)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_progress NUMERIC(5,2)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_run_id TEXT" in sql
assert "ADD COLUMN IF NOT EXISTS preview_failure_reason TEXT" in sql
assert "ADD COLUMN IF NOT EXISTS preview_total_files INTEGER" in sql
assert "ADD COLUMN IF NOT EXISTS preview_completed_files INTEGER" in sql
assert "data_process_workflow_backfill_ids" in sql
assert "ck_data_process_tasks_workflow_step" in sql
assert "ck_data_process_tasks_preview_status" in sql
assert "ck_data_process_tasks_preview_progress" in sql
assert "ck_data_process_tasks_preview_file_counts" in sql
for value in ("create", "model", "upload", "preview", "generate", "results"):
assert f"'{value}'" in sql
for value in ("idle", "queued", "running", "completed", "failed", "cancelled"):
assert f"'{value}'" in sql
assert "CREATE TABLE IF NOT EXISTS data_process_results" in sql
assert sql.count("BEGIN;") == 1
assert sql.rstrip().endswith("COMMIT;")
@@ -27,3 +45,17 @@ def test_schema_cli_target_label_never_contains_credentials() -> None:
label = _target_label("postgresql://secret-user:secret-password@db.example:5433/yg_ft")
assert label == "db.example:5433/yg_ft"
assert "secret" not in label
def test_schema_check_requires_current_runtime_columns() -> None:
assert REQUIRED_TASK_COLUMNS == (
"generation_run_id",
"results_confirmed",
"workflow_step",
"preview_status",
"preview_progress",
"preview_run_id",
"preview_failure_reason",
"preview_total_files",
"preview_completed_files",
)

111
backend/tests/test_data_process_store.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import inspect
import json
from collections.abc import Iterator
from contextlib import contextmanager
@@ -20,6 +21,14 @@ from app.modules.data_process.store import (
)
def test_preview_replace_sql_never_uses_untyped_null_placeholders() -> None:
source = inspect.getsource(DataProcessStore.replace_preview_items)
assert "%s IS NULL" not in source
assert "is_direct_build = preview_run_id is None" in source
assert "workflow_step=CASE WHEN %s THEN 'preview'" in source
class _Result:
def __init__(self, *, row: dict[str, Any] | None = None, rows: list[dict[str, Any]] | None = None):
self.row = row
@@ -146,6 +155,7 @@ class _PublishStore(DataProcessStore):
return {
"id": task_id,
"status": "completed",
"results_confirmed": True,
"description": "",
"config": self._task_config,
"output_dataset_id": train_dataset and train_dataset["id"],
@@ -454,12 +464,14 @@ class _LegacyRecoveryConnection:
record["preview_item_id"] = params[1]
return _Result()
if normalized.startswith("UPDATE data_process_tasks SET status='completed'"):
assert "workflow_step='results'" in normalized
self.task.update(
{
"status": "completed",
"progress": 100,
"output_dataset_id": params[0],
"output_count": params[1],
"workflow_step": "results",
}
)
return _Result()
@@ -496,6 +508,7 @@ class _StartGenerationConnection:
config=config,
output_dataset_id="dataset_train" if published_prepared else None,
output_count=28,
results_confirmed=published_prepared,
),
"generation_run_id": None,
}
@@ -512,6 +525,7 @@ class _StartGenerationConnection:
assert normalized.startswith("UPDATE data_process_tasks SET config=%s, status='running'")
assert "output_dataset_id=NULL" in normalized
assert "output_count=0" in normalized
assert "results_confirmed=FALSE" in normalized
self.task.update(
{
"config": params[0],
@@ -526,6 +540,7 @@ class _StartGenerationConnection:
"duplicate_count": 0,
"error_count": 0,
"generation_run_id": params[2],
"results_confirmed": False,
"updated_at": params[3],
}
)
@@ -713,6 +728,7 @@ def test_start_generation_clears_previous_output_count() -> None:
assert task["status"] == "running"
assert task["output_count"] == 0
assert task["results_confirmed"] is False
assert conn.results == []
@@ -737,6 +753,7 @@ def test_prepared_published_task_survives_generation_preflight_failure() -> None
assert conn.task["status"] == "completed"
assert conn.task["output_dataset_id"] == "dataset_train"
assert conn.task["output_count"] == 28
assert conn.task["results_confirmed"] is True
assert "_regeneration_prepared" in conn.task["config"]
assert conn.results == [{"id": "old-result"}]
@@ -752,6 +769,7 @@ def test_legacy_aborted_regeneration_recovers_results_and_published_state() -> N
assert conn.task["progress"] == 100
assert conn.task["output_dataset_id"] == "dataset_train"
assert conn.task["output_count"] == 2
assert conn.task["workflow_step"] == "results"
assert conn.task["started_at"] is None
assert conn.task["completed_at"] is None
assert [item["id"] for item in conn.results] == ["result_train", "result_test"]
@@ -1185,3 +1203,96 @@ def test_source_storage_descriptor_rejects_unowned_or_unsupported_references(
"dpt_task",
"dpsf_source",
)
class _LifecycleConnection:
def __init__(self) -> None:
self.task: dict[str, Any] = {
"id": "task-lifecycle",
"status": "running",
"generation_run_id": "generation-active",
"workflow_step": "generate",
"preview_status": "running",
"preview_progress": Decimal("40.00"),
"preview_run_id": "preview-active",
"preview_failure_reason": None,
"preview_total_files": 5,
"preview_completed_files": 2,
"deleted_at": None,
"deleted_by": None,
}
self.last_update_sql = ""
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("SELECT * FROM data_process_tasks"):
row = None if self.task["deleted_at"] is not None else dict(self.task)
return _Result(row=row)
if normalized.startswith("UPDATE data_process_tasks SET workflow_step="):
self.last_update_sql = normalized
workflow_step, updated_at, task_id = params
assert task_id == self.task["id"]
self.task.update(workflow_step=workflow_step, updated_at=updated_at)
return _Result(row=dict(self.task))
if normalized.startswith("UPDATE data_process_tasks SET status=CASE"):
self.last_update_sql = normalized
deleted_at, deleted_by, updated_at, task_id = params
assert task_id == self.task["id"]
self.task.update(
status="stopped",
generation_run_id=None,
preview_status="cancelled",
preview_run_id=None,
deleted_at=deleted_at,
deleted_by=deleted_by,
updated_at=updated_at,
)
return _Result()
if normalized.startswith("SELECT status, generation_run_id"):
row = None if self.task["deleted_at"] is not None else dict(self.task)
return _Result(row=row)
if normalized.startswith("SELECT preview_status, preview_run_id"):
row = None if self.task["deleted_at"] is not None else dict(self.task)
return _Result(row=row)
raise AssertionError(f"unexpected SQL: {normalized}")
class _LifecycleStore(DataProcessStore):
def __init__(self, conn: _LifecycleConnection) -> None:
self._conn = conn
@contextmanager
def connect(self) -> Iterator[_LifecycleConnection]:
yield self._conn
def test_workflow_step_update_does_not_invalidate_active_runs() -> None:
conn = _LifecycleConnection()
task = _LifecycleStore(conn).update_workflow_step("task-lifecycle", "results")
assert task["workflow_step"] == "results"
assert task["status"] == "running"
assert task["generation_run_id"] == "generation-active"
assert task["preview_status"] == "running"
assert task["preview_run_id"] == "preview-active"
assert "generation_run_id" not in conn.last_update_sql
assert "preview_run_id" not in conn.last_update_sql
def test_delete_atomically_invalidates_generation_and_preview_runs() -> None:
conn = _LifecycleConnection()
store = _LifecycleStore(conn)
store.delete_task("task-lifecycle", deleted_by="user-1")
assert conn.task["status"] == "stopped"
assert conn.task["generation_run_id"] is None
assert conn.task["preview_status"] == "cancelled"
assert conn.task["preview_run_id"] is None
assert conn.task["deleted_by"] == "user-1"
assert conn.task["deleted_at"] is not None
assert store.generation_is_running("task-lifecycle", "generation-active") is False
assert store.preview_is_running("task-lifecycle", "preview-active") is False