YG_FT/backend/tests/test_data_process_migration.py

from __future__ import annotations

from pathlib import Path

from app.modules.data_process.schema_cli import REQUIRED_TASK_COLUMNS, _target_label


def test_runtime_migration_fails_fast_on_incompatible_schema() -> None:
    sql_path = (
        Path(__file__).resolve().parents[1]
        / "app"
        / "db"
        / "sql"
        / "002_data_process.sql"
    )
    sql = sql_path.read_text(encoding="utf-8")

    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 "ADD COLUMN IF NOT EXISTS chosen TEXT NOT NULL DEFAULT ''" in sql
    assert "ADD COLUMN IF NOT EXISTS rejected TEXT NOT NULL DEFAULT ''" in sql
    assert "ADD COLUMN IF NOT EXISTS original_chosen TEXT" in sql
    assert "ADD COLUMN IF NOT EXISTS original_rejected TEXT" in sql
    assert sql.count("BEGIN;") == 1
    assert sql.rstrip().endswith("COMMIT;")


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