stack/tests/aco/test_table_base.py

"""Tests for aco.table.base — SQLTable base class."""

from __future__ import annotations

from datetime import date, datetime
from decimal import Decimal

from pydantic import Field

from aco.table.base import SQLTable, _resolve_type


class TestSQLTable:
    def test_qualified_name(self) -> None:
        class MyTable(SQLTable):
            __schema__ = "core"
            __tablename__ = "encounter"
            col_a: str | None = None

        assert MyTable.qualified_name() == "core.encounter"

    def test_column_names(self) -> None:
        class MyTable(SQLTable):
            __schema__ = "test"
            __tablename__ = "demo"
            col_a: str | None = None
            col_b: int | None = None

        assert MyTable.column_names() == ["col_a", "col_b"]

    def test_qualified_name_with_real_model(self) -> None:
        from aco.table.alr import AlrAssignedBeneficiaries

        assert (
            AlrAssignedBeneficiaries.qualified_name()
            == "alr.alr_assigned_beneficiaries"
        )

    def test_column_names_with_real_model(self) -> None:
        from aco.table.alr import AlrAssignedBeneficiaries

        names = AlrAssignedBeneficiaries.column_names()
        assert "mbi" in names
        assert len(names) > 10


class TestColumnMeta:
    def test_column_meta_with_descriptions(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str = Field(description="first col")
            b: int | None = None

        meta = T.column_meta()
        assert meta["a"]["description"] == "first col"
        assert meta["b"]["description"] == ""

    def test_column_meta_with_json_schema_extra(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str = Field(
                description="id",
                json_schema_extra={"source": "input_layer.claim"},
            )

        meta = T.column_meta()
        assert meta["a"]["extra"]["source"] == "input_layer.claim"

    def test_column_meta_empty_extra_when_unset(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str | None = None

        meta = T.column_meta()
        assert meta["a"]["extra"] == {}

    def test_column_meta_includes_type(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: int | None = None

        meta = T.column_meta()
        assert meta["a"]["type"] is not None


class TestFieldDescriptions:
    def test_returns_only_described_fields(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str = Field(description="has desc")
            b: int | None = None
            c: str | None = Field(default=None, description="also has desc")

        descs = T.field_descriptions()
        assert descs == {"a": "has desc", "c": "also has desc"}

    def test_empty_when_no_descriptions(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str | None = None

        assert T.field_descriptions() == {}


class TestResolveType:
    def test_str(self) -> None:
        assert _resolve_type(str) == "VARCHAR"

    def test_int(self) -> None:
        assert _resolve_type(int) == "INTEGER"

    def test_float(self) -> None:
        assert _resolve_type(float) == "DOUBLE"

    def test_bool(self) -> None:
        assert _resolve_type(bool) == "BOOLEAN"

    def test_date(self) -> None:
        assert _resolve_type(date) == "DATE"

    def test_datetime(self) -> None:
        assert _resolve_type(datetime) == "TIMESTAMP"

    def test_decimal(self) -> None:
        assert _resolve_type(Decimal) == "DECIMAL(18,2)"

    def test_nullable_str(self) -> None:
        assert _resolve_type(str | None) == "VARCHAR"

    def test_nullable_int(self) -> None:
        assert _resolve_type(int | None) == "INTEGER"

    def test_unknown_falls_back_to_varchar(self) -> None:
        assert _resolve_type(object) == "VARCHAR"


class TestToDdl:
    def test_basic_ddl(self) -> None:
        class T(SQLTable):
            __schema__ = "core"
            __tablename__ = "demo"
            a: str | None = None
            b: int | None = None

        ddl = T.to_ddl()
        assert ddl.startswith("CREATE TABLE core.demo")
        assert "a VARCHAR" in ddl
        assert "b INTEGER" in ddl

    def test_ddl_with_descriptions_as_comments(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str = Field(description="the a column")
            b: int | None = None

        ddl = T.to_ddl()
        assert "-- the a column" in ddl
        assert "b INTEGER" in ddl

    def test_ddl_comma_placement(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str | None = None
            b: int | None = None

        ddl = T.to_ddl()
        lines = ddl.strip().split("\n")
        # first col line should have comma, last should not
        assert lines[1].strip().startswith("a VARCHAR,")
        assert not lines[2].strip().endswith(",")

    def test_ddl_all_types(self) -> None:
        class T(SQLTable):
            __schema__ = "s"
            __tablename__ = "t"
            a: str | None = None
            b: int | None = None
            c: float | None = None
            d: bool | None = None
            e: date | None = None
            f: datetime | None = None
            g: Decimal | None = None

        ddl = T.to_ddl()
        assert "VARCHAR" in ddl
        assert "INTEGER" in ddl
        assert "DOUBLE" in ddl
        assert "BOOLEAN" in ddl
        assert "DATE" in ddl
        assert "TIMESTAMP" in ddl
        assert "DECIMAL" in ddl