tests/conftest.py

"""Test-wide fixtures.

Two complementary fixture families live here:

  **PKDD-backed fixtures (Bhavish)** — for B-owned tools that need the
  bundled DuckDB and optionally a real LLM:
    * `pkdd_handle` / `pkdd_ctx` — depend on the DuckDB being bootstrapped;
      `pytest` skips cleanly if it's missing.
    * `lexsi_env` / `pkdd_ctx_lexsi` — gated behind `SDK_ACCESS_TOKEN` +
      project env vars; skip cleanly when they're absent.

  **Offline-synthetic fixtures (Aditya)** — for the modeling tools
  (`train_tabular_model`, `predict`, `explain_prediction`); fully
  hermetic, no DuckDB / Lexsi / LLM dependency:
    * `loan_context_df` / `loan_predict_df` — small PKDD-loan-shaped
      DataFrames with / without labels.
    * `fake_dataset_handle`, `fake_tab_project`, `ctx_with_fake`,
      `ctx_offline` — context wired against `FakeTabularProject` from
      `lexsi_ds.agent.testing.fakes`.
    * `seed_sql_result` — helper to stash a DataFrame on `ctx.cache` the
      way `run_sql` would.

No fixture-name collisions; the two families are pytest-discovered
together.
"""

from __future__ import annotations

import os
import random
from pathlib import Path

import pandas as pd
import pytest

from lexsi_ds.agent.context import AgentContext, DatasetHandle
from lexsi_ds.agent.testing import FakeTabularProject
from lexsi_ds.paths import DUCKDB_PATH


# ===========================================================================
# PKDD-backed fixtures 
# ===========================================================================


@pytest.fixture(scope="session")
def pkdd_handle():
    """Load the bundled PKDD DatasetHandle. Skip if the DuckDB isn't built yet."""
    if not DUCKDB_PATH.exists():
        pytest.skip(
            f"PKDD DuckDB not built at {DUCKDB_PATH}. Run "
            f"`uv run python -m lexsi_ds.data.loader` first."
        )
    from lexsi_ds.agent.datasource import load_pkdd_handle

    return load_pkdd_handle()


@pytest.fixture()
def pkdd_ctx(pkdd_handle):
    """Minimal AgentContext on PKDD with the stub LLM. Good for non-LLM tools."""
    from lexsi_ds.agent.context import AgentContext
    from lexsi_ds.llm.client import factory as llm_factory

    return AgentContext(
        dataset=pkdd_handle,
        run_id="test",
        llm=llm_factory("stub"),
    )


@pytest.fixture()
def lexsi_env() -> dict[str, str]:
    """Skip the test unless the Lexsi env is fully set."""
    required = ("SDK_ACCESS_TOKEN", "LEXSI_WORKSPACE_NAME", "LEXSI_TEXT_PROJECT_NAME")
    missing = [v for v in required if not os.environ.get(v)]
    if missing:
        pytest.skip(
            f"Lexsi env not configured (missing: {missing}). "
            f"See README quickstart for setup."
        )
    return {v: os.environ[v] for v in required}


@pytest.fixture()
def pkdd_ctx_lexsi(pkdd_handle, lexsi_env):
    """AgentContext on PKDD with a real LexsiTextClient."""
    from lexsi_ds.agent.context import AgentContext
    from lexsi_ds.llm.client import factory as llm_factory

    return AgentContext(
        dataset=pkdd_handle,
        run_id="test-lexsi",
        llm=llm_factory("lexsi"),
    )


# ===========================================================================
# Offline-synthetic fixtures
# ===========================================================================


@pytest.fixture
def loan_context_df() -> pd.DataFrame:
    """Tiny PKDD-loan-default-shaped DataFrame with KNOWN labels (context)."""
    rng = random.Random(42)
    n = 60
    rows = []
    for i in range(n):
        amount = rng.randint(5_000, 200_000)
        duration = rng.choice([12, 24, 36, 48, 60])
        payments = round(amount / duration * rng.uniform(0.9, 1.1), 2)
        # Default risk roughly proportional to amount / duration.
        risk = (amount / 200_000) * (1 - duration / 60)
        status = "defaulted" if rng.random() < risk else "repaid"
        rows.append({
            "loan_id": 1000 + i,
            "amount": amount,
            "duration": duration,
            "payments": payments,
            "status": status,
        })
    return pd.DataFrame(rows)


@pytest.fixture
def loan_predict_df() -> pd.DataFrame:
    """Predict-set: same shape, no status column (label-unknown)."""
    rng = random.Random(7)
    n = 25
    rows = []
    for i in range(n):
        amount = rng.randint(5_000, 200_000)
        duration = rng.choice([12, 24, 36, 48, 60])
        payments = round(amount / duration * rng.uniform(0.9, 1.1), 2)
        rows.append({
            "loan_id": 9000 + i,
            "amount": amount,
            "duration": duration,
            "payments": payments,
            "status": None,           # all-null target — predict tool will strip
        })
    return pd.DataFrame(rows)


@pytest.fixture
def fake_dataset_handle(tmp_path: Path) -> DatasetHandle:
    """A minimal in-memory DatasetHandle. We don't touch DuckDB in tests."""
    return DatasetHandle(
        id="test_offline",
        kind="bundled",
        duckdb_path=tmp_path / "nonexistent.duckdb",
        tables=[],
        connector_id=None,
        kg=None,
    )


@pytest.fixture
def fake_tab_project() -> FakeTabularProject:
    return FakeTabularProject(project_name="test_proj", seed=42)


@pytest.fixture
def ctx_with_fake(
    fake_dataset_handle: DatasetHandle,
    fake_tab_project: FakeTabularProject,
) -> AgentContext:
    """A test context with a FakeTabularProject wired in (no LLM)."""
    return AgentContext(
        dataset=fake_dataset_handle,
        run_id="test_run",
        org=None,
        text_project=None,
        tab_project=fake_tab_project,
        llm=None,
    )


@pytest.fixture
def ctx_offline(fake_dataset_handle: DatasetHandle) -> AgentContext:
    """A test context with NO tab_project — for negative-path tests."""
    return AgentContext(
        dataset=fake_dataset_handle,
        run_id="test_offline",
        org=None,
        text_project=None,
        tab_project=None,
        llm=None,
    )


def _seed_sql_result(ctx: AgentContext, label: str, df: pd.DataFrame) -> None:
    """Helper: stash a DataFrame on ctx.cache the way run_sql would."""
    ctx.cache[f"sql_result:{label}"] = df
    ctx.cache["last_sql_result"] = df


@pytest.fixture
def seed_sql_result():
    return _seed_sql_result