tests/test_sample_values_idents.py

"""sample_values must handle real-world column names with spaces/punctuation
('Listing Exchange', 'Company Description') — the #1 error in the DAB
trajectory analysis (11x bad_identifier). Double-quoting makes them safe."""

from __future__ import annotations

from pathlib import Path

import duckdb

from lexsi_ds.agent.context import AgentContext, ColumnInfo, DatasetHandle, TableInfo
from lexsi_ds.agent.tools.sample_values import SampleValuesArgs, _run


def _ctx(tmp_path):
    p = tmp_path / "d.duckdb"
    con = duckdb.connect(str(p))
    con.execute('CREATE TABLE stockinfo ("Listing Exchange" VARCHAR, "Company Description" VARCHAR)')
    con.executemany('INSERT INTO stockinfo VALUES (?, ?)',
                    [("N", "Apple"), ("N", "MFA"), ("Q", "Goog")])
    con.close()
    ds = DatasetHandle(id="t", kind="attached", duckdb_path=p, tables=[
        TableInfo(name="stockinfo", columns=[
            ColumnInfo("Listing Exchange", "VARCHAR"),
            ColumnInfo("Company Description", "VARCHAR")])])
    return AgentContext(dataset=ds, run_id="t")


def test_spaced_column_name_works(tmp_path):
    res = _run(SampleValuesArgs(table="stockinfo", column="Listing Exchange"), _ctx(tmp_path))
    assert res.ok, res.summary
    assert set(res.payload["values"]) == {"N", "Q"}


def test_already_quoted_column_name_unwrapped(tmp_path):
    res = _run(SampleValuesArgs(table="stockinfo", column='"Listing Exchange"'), _ctx(tmp_path))
    assert res.ok, res.summary
    assert set(res.payload["values"]) == {"N", "Q"}


def test_truly_unsafe_identifier_still_rejected(tmp_path):
    res = _run(SampleValuesArgs(table="stockinfo", column="bad\nname"), _ctx(tmp_path))
    assert not res.ok and res.error == "bad_identifier"