test_cli.py

import json
from unittest.mock import patch

from test_croissant import CROISSANT_FIXTURE
from typer.testing import CliRunner

from otai import catalog, sql_guard
from otai.cli import app

runner = CliRunner()

SAMPLE_LISTING_XML = b"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
    <CommonPrefixes><Prefix>platform/25.12/</Prefix></CommonPrefixes>
    <CommonPrefixes><Prefix>platform/26.06/</Prefix></CommonPrefixes>
</ListBucketResult>
"""


def _invoke(args, cache_dir):
    with patch(
        "otai.releases.default_fetch_listing_xml", return_value=SAMPLE_LISTING_XML
    ):
        return runner.invoke(app, args, env={"OTAI_CACHE_DIR": str(cache_dir)})


def _invoke_with_fixtures(args, cache_dir, base_uri):
    with (
        patch(
            "otai.releases.default_fetch_listing_xml", return_value=SAMPLE_LISTING_XML
        ),
        patch(
            "otai.croissant.default_fetch_croissant",
            return_value=json.dumps(CROISSANT_FIXTURE).encode(),
        ),
    ):
        return runner.invoke(
            app,
            args,
            env={"OTAI_CACHE_DIR": str(cache_dir), "OTAI_BASE_URI": base_uri},
        )


def test_list_releases_json_output_default(tmp_path):
    result = _invoke(["list-releases"], tmp_path / "cache")

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is True
    releases = {r["release"]: r for r in payload["data"]["releases"]}
    assert releases["26.06"]["latest"] is True
    assert releases["25.12"]["latest"] is False


def test_list_releases_table_format(tmp_path):
    result = _invoke(["list-releases", "--format", "table"], tmp_path / "cache")

    assert result.exit_code == 0
    assert "release" in result.stdout
    assert "26.06" in result.stdout
    assert "yes" in result.stdout


def test_list_releases_has_no_release_flag(tmp_path):
    result = _invoke(["list-releases", "--release", "26.06"], tmp_path / "cache")
    assert result.exit_code != 0


def test_list_releases_does_not_create_catalog_file(tmp_path):
    # list-releases never builds a schema, so it must never take the
    # catalog's write lock - it should simply report nothing cached rather
    # than unconditionally creating the file.
    cache_dir = tmp_path / "cache"
    assert not (cache_dir / "catalog.duckdb").exists()

    result = _invoke(["list-releases"], cache_dir)

    assert result.exit_code == 0
    assert not (cache_dir / "catalog.duckdb").exists()


def test_unknown_format_is_rejected(tmp_path):
    result = _invoke(["list-releases", "--format", "yaml"], tmp_path / "cache")
    assert result.exit_code != 0


def test_list_datasets_json_output_defaults_to_latest(tmp_path, fixture_release_layout):
    base_uri, release, dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(["list-datasets"], tmp_path / "cache", base_uri)

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is True
    assert payload["data"]["release"] == release
    names = {d["dataset"] for d in payload["data"]["datasets"]}
    assert names == set(dataset_rows)


def test_list_datasets_table_format(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["list-datasets", "--format", "table"], tmp_path / "cache", base_uri
    )

    assert result.exit_code == 0
    assert "dataset" in result.stdout
    assert "description" in result.stdout
    assert "target" in result.stdout


def test_list_datasets_accepts_explicit_release(tmp_path, fixture_release_layout):
    base_uri, release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["list-datasets", "--release", release], tmp_path / "cache", base_uri
    )

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["data"]["release"] == release


def test_describe_dataset_json_output_defaults_to_latest(
    tmp_path, fixture_release_layout
):
    base_uri, release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["describe-dataset", "target"], tmp_path / "cache", base_uri
    )

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is True
    assert payload["data"]["release"] == release
    assert payload["data"]["dataset"] == "target"
    fields_by_name = {f["name"]: f for f in payload["data"]["fields"]}
    assert fields_by_name["id"]["dataType"] == "sc:Text"


def test_describe_dataset_includes_references_and_subfields(
    tmp_path, fixture_release_layout
):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["describe-dataset", "association_by_datasource_direct"],
        tmp_path / "cache",
        base_uri,
    )

    payload = json.loads(result.stdout)
    fields_by_name = {f["name"]: f for f in payload["data"]["fields"]}
    assert fields_by_name["targetId"]["references"] == {
        "dataset": "target",
        "field": "id",
    }

    result = _invoke_with_fixtures(
        ["describe-dataset", "target"], tmp_path / "cache", base_uri
    )
    payload = json.loads(result.stdout)
    fields_by_name = {f["name"]: f for f in payload["data"]["fields"]}
    sub_by_name = {sf["name"]: sf for sf in fields_by_name["proteinIds"]["subFields"]}
    assert set(sub_by_name) == {"id", "source"}


def test_describe_dataset_table_format(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["describe-dataset", "target", "--format", "table"],
        tmp_path / "cache",
        base_uri,
    )

    assert result.exit_code == 0
    assert "name" in result.stdout
    assert "dataType" in result.stdout
    assert "id" in result.stdout


def test_describe_dataset_accepts_explicit_release(tmp_path, fixture_release_layout):
    base_uri, release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["describe-dataset", "target", "--release", release],
        tmp_path / "cache",
        base_uri,
    )

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["data"]["release"] == release


def test_describe_dataset_unknown_dataset_returns_error(
    tmp_path, fixture_release_layout
):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["describe-dataset", "no_such_dataset"], tmp_path / "cache", base_uri
    )

    assert result.exit_code != 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is False
    assert payload["error"]["type"] == "dataset_not_found"


def test_list_datasets_builds_catalog_schema_on_first_run(
    tmp_path, fixture_release_layout
):
    base_uri, release, _dataset_rows = fixture_release_layout
    cache_dir = tmp_path / "cache"

    _invoke_with_fixtures(["list-datasets"], cache_dir, base_uri)

    import duckdb

    conn = duckdb.connect(str(cache_dir / "catalog.duckdb"))
    try:
        schemas = {
            row[0]
            for row in conn.execute(
                "SELECT schema_name FROM information_schema.schemata"
            ).fetchall()
        }
        assert release in schemas
    finally:
        conn.close()


def test_run_sql_json_output_against_latest(tmp_path, fixture_release_layout):
    base_uri, release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["run-sql", "SELECT id, approvedSymbol FROM target ORDER BY id"],
        tmp_path / "cache",
        base_uri,
    )

    assert result.exit_code == 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is True
    assert payload["data"]["release"] == release
    assert payload["data"]["columns"] == ["id", "approvedSymbol"]
    assert payload["data"]["rows"] == [
        ["ENSG00000141510", "TP53"],
        ["ENSG00000157764", "BRAF"],
    ]


def test_run_sql_table_format(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        [
            "run-sql",
            "SELECT id, approvedSymbol FROM target ORDER BY id",
            "--format",
            "table",
        ],
        tmp_path / "cache",
        base_uri,
    )

    assert result.exit_code == 0
    assert "approvedSymbol" in result.stdout
    assert "BRAF" in result.stdout


def test_run_sql_has_no_release_flag(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["run-sql", "SELECT 1", "--release", "26.06"], tmp_path / "cache", base_uri
    )

    assert result.exit_code != 0


def test_run_sql_rejects_non_select_statement(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["run-sql", "DROP TABLE target"], tmp_path / "cache", base_uri
    )

    assert result.exit_code != 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is False
    assert payload["error"]["type"] == "guardrail_violation"


def test_run_sql_malformed_query_returns_sql_error(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["run-sql", "SELECT * FROM target WHERE ("], tmp_path / "cache", base_uri
    )

    assert result.exit_code != 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is False
    assert payload["error"]["type"] == "sql_error"


def test_run_sql_rejects_non_positive_timeout(tmp_path, fixture_release_layout):
    base_uri, _release, _dataset_rows = fixture_release_layout

    result = _invoke_with_fixtures(
        ["run-sql", "SELECT 1", "--timeout", "0"], tmp_path / "cache", base_uri
    )

    # Match this file's convention for invalid-option tests (e.g.
    # test_unknown_format_is_rejected): check the exit code only, not the
    # styled error text - Typer/Click's box-drawn error panel wraps
    # differently depending on the terminal width Rich detects, so exact
    # substring matches against result.output are environment-fragile.
    assert result.exit_code == 2


def test_run_sql_passes_custom_timeout_through_to_commands(
    tmp_path, fixture_release_layout
):
    base_uri, _release, _dataset_rows = fixture_release_layout

    with patch("otai.cli.commands.run_sql") as mock_run_sql:
        mock_run_sql.return_value = {"ok": True, "data": {"columns": [], "rows": []}}
        _invoke_with_fixtures(
            ["run-sql", "SELECT 1", "--timeout", "120"], tmp_path / "cache", base_uri
        )

    assert mock_run_sql.call_args.kwargs["timeout_seconds"] == 120.0


def test_run_sql_uses_default_timeout_when_flag_omitted(
    tmp_path, fixture_release_layout
):
    base_uri, _release, _dataset_rows = fixture_release_layout

    with patch("otai.cli.commands.run_sql") as mock_run_sql:
        mock_run_sql.return_value = {"ok": True, "data": {"columns": [], "rows": []}}
        _invoke_with_fixtures(["run-sql", "SELECT 1"], tmp_path / "cache", base_uri)

    assert (
        mock_run_sql.call_args.kwargs["timeout_seconds"]
        == sql_guard.DEFAULT_TIMEOUT_SECONDS
    )


def test_run_sql_timeout_flag_actually_shortens_execution(
    tmp_path, fixture_release_layout
):
    # Real end-to-end confirmation (not mocked): a view wrapping range()
    # (created outside the guard, same pattern used in test_commands.py/
    # test_sql_guard.py) is cheap to set up but genuinely slow to query -
    # a short --timeout must cause it to time out for real.
    base_uri, release, _dataset_rows = fixture_release_layout
    cache_dir = tmp_path / "cache"

    _invoke_with_fixtures(["list-datasets"], cache_dir, base_uri)
    conn = catalog.connect_catalog(cache_dir)
    try:
        conn.execute(
            f'CREATE VIEW "{release}".slow_a AS SELECT * FROM range(100000000)'
        )
        conn.execute(f'CREATE VIEW "{release}".slow_b AS SELECT * FROM range(100000)')
    finally:
        conn.close()

    result = _invoke_with_fixtures(
        ["run-sql", "SELECT count(*) FROM slow_a a, slow_b b", "--timeout", "0.2"],
        cache_dir,
        base_uri,
    )

    assert result.exit_code != 0
    payload = json.loads(result.stdout)
    assert payload["ok"] is False
    assert payload["error"]["type"] == "timeout"
    assert "0.2" in payload["error"]["message"]