archon/tests/test_sql.py at master · twomathematicians-code/archon · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
"""Tests for SQL backend code generation."""

from archon.codegen.backends.sql import (
    generate_insert_template,
    generate_select,
    generate_table_ddl,
)
from archon.dsl.schema import DataDomain, DataFeature, DataSchema


def make_schema(name="test_table", features=None, target=None):
    if features is None:
        features = [
            DataFeature(name="id", dtype="int32"),
            DataFeature(name="name", dtype="str"),
            DataFeature(name="score", dtype="float64"),
        ]
    return DataSchema(
        name=name,
        domain=DataDomain.TABULAR,
        source="./data",
        format="csv",
        features=features,
        target=target,
        seed=42,
    )


def test_generate_table_ddl():
    schema = make_schema()
    ddl = generate_table_ddl(schema)
    assert "CREATE TABLE test_table" in ddl
    assert "id INTEGER" in ddl
    assert "name VARCHAR(255)" in ddl
    assert "score DOUBLE PRECISION" in ddl


def test_generate_table_ddl_not_null():
    schema = make_schema(
        features=[
            DataFeature(name="id", dtype="int32", nullable=False),
        ]
    )
    ddl = generate_table_ddl(schema)
    assert "NOT NULL" in ddl


def test_generate_table_ddl_nullable():
    schema = make_schema(
        features=[
            DataFeature(name="opt", dtype="str", nullable=True),
        ]
    )
    ddl = generate_table_ddl(schema)
    # Nullable columns should not have NOT NULL
    assert "NOT NULL" not in ddl


def test_generate_table_ddl_with_target():
    schema = make_schema(target="score")
    ddl = generate_table_ddl(schema)
    assert "-- target column: score" in ddl


def test_generate_insert_template():
    schema = make_schema()
    insert = generate_insert_template(schema)
    assert "INSERT INTO test_table" in insert
    assert "(id, name, score)" in insert
    assert "VALUES (?, ?, ?)" in insert


def test_generate_select():
    schema = make_schema()
    select = generate_select(schema, limit=50)
    assert "SELECT id, name, score" in select
    assert "FROM test_table" in select
    assert "LIMIT 50" in select


def test_generate_select_default_limit():
    schema = make_schema(features=[DataFeature(name="x", dtype="float64")])
    select = generate_select(schema)
    assert "LIMIT 100" in select


def test_type_mapping():
    schema = make_schema(
        features=[
            DataFeature(name="a", dtype="float32"),
            DataFeature(name="b", dtype="int64"),
            DataFeature(name="c", dtype="bool"),
            DataFeature(name="d", dtype="datetime"),
        ]
    )
    ddl = generate_table_ddl(schema)
    assert "a FLOAT" in ddl
    assert "b BIGINT" in ddl
    assert "c BOOLEAN" in ddl
    assert "d TIMESTAMP" in ddl


def test_type_mapping_defaults():
    """Verify dtype 'str' maps to VARCHAR."""
    schema = make_schema(
        features=[
            DataFeature(name="name", dtype="str"),
        ]
    )
    ddl = generate_table_ddl(schema)
    assert "name VARCHAR(255)" in ddl