tests/test_smoke.py

"""Fast smoke tests for CI (DuckDB warm-up once per process)."""

from __future__ import annotations

import os
import sys

import pytest

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from env import SQLOptimEnv  # noqa: E402
from graders import GradeMask, grade  # noqa: E402
from models import Action  # noqa: E402
from tasks import TASKS  # noqa: E402


@pytest.fixture(scope="module")
def executor():
    from executor import get_executor

    return get_executor()


def test_executor_compare_task1(executor):
    task = TASKS["task_1_basic_antipatterns"]
    original = task["sql_query"]
    optimized = (
        "SELECT id, customer_id, product_id, status, total, created_at FROM orders "
        "WHERE customer_id = 5000 "
        "AND created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01'"
    )
    r = executor.compare(original, optimized)
    assert r["speedup"] >= 1.0
    assert r["results_match"] is True


def test_grade_mask_changes_total():
    task = TASKS["task_1_basic_antipatterns"]
    action = Action(
        suggestions=[
            {
                "issue_type": "select_star",
                "line": 1,
                "description": "SELECT * on large table",
                "severity": "high",
                "fix": "project columns",
            }
        ],
        optimized_query=(
            "SELECT id, customer_id, product_id, status, total, created_at FROM orders "
            "WHERE customer_id = 5000 AND created_at >= DATE '2024-01-01' "
            "AND created_at < DATE '2025-01-01'"
        ),
        summary="x" * 130,
        estimated_improvement="5x",
        approved=False,
    )
    full = grade(task, action).score
    no_exec = grade(
        task, action, mask=GradeMask(execution_speedup=False, result_correctness=False)
    ).score
    assert no_exec < full


def test_fastapi_reset_step():
    from fastapi.testclient import TestClient

    from server.app import app

    client = TestClient(app)
    r = client.get("/")
    assert r.status_code == 200
    assert r.json()["environment"] == "sql-optim-env"

    obs = client.post("/reset", json={"task_id": "task_1_basic_antipatterns"}).json()
    assert obs["task_id"] == "task_1_basic_antipatterns"

    step = client.post(
        "/step",
        json={
            "suggestions": [
                {
                    "issue_type": "select_star",
                    "line": 1,
                    "description": "avoid star",
                    "severity": "high",
                    "fix": "cols",
                }
            ],
            "optimized_query": (
                "SELECT id, customer_id, product_id, status, total, created_at FROM orders "
                "WHERE customer_id = 5000 "
                "AND created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01'"
            ),
            "summary": "Rewrite removes anti-patterns and uses a sargable date range.",
            "estimated_improvement": "4x",
            "approved": False,
        },
    )
    assert step.status_code == 200
    body = step.json()
    assert "reward" in body
    assert body["reward"]["score"] > 0.5


def test_sqoptim_env_reset_step():
    env = SQLOptimEnv()
    obs = env.reset(task_id="task_1_basic_antipatterns")
    assert obs.step_count == 0
    result = env.step(
        Action(
            suggestions=[
                {
                    "issue_type": "select_star",
                    "line": 1,
                    "description": "SELECT *",
                    "severity": "high",
                    "fix": "list columns",
                }
            ],
            optimized_query=(
                "SELECT id, customer_id, product_id, status, total, created_at FROM orders "
                "WHERE customer_id = 5000 "
                "AND created_at >= DATE '2024-01-01' AND created_at < DATE '2025-01-01'"
            ),
            summary="A" * 130,
            estimated_improvement="5x",
            approved=False,
        )
    )
    assert result.reward.score > 0.4