sqlmesh/tests/utils/test_helpers.py at 1fb73e4d14849c601c4c65a58896c7d8c544bf6f · SQLMesh/sqlmesh · 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
import pytest
from functools import wraps
from sqlglot import expressions
from sqlglot.optimizer.annotate_types import annotate_types

from sqlmesh.core.console import set_console, get_console, TerminalConsole

from sqlmesh.utils import columns_to_types_all_known


@pytest.mark.parametrize(
    "columns_to_types, expected",
    [
        ({"a": expressions.DataType.build("INT"), "b": expressions.DataType.build("INT")}, True),
        (
            {"a": expressions.DataType.build("UNKNOWN"), "b": expressions.DataType.build("INT")},
            False,
        ),
        ({"a": expressions.DataType.build("NULL"), "b": expressions.DataType.build("INT")}, False),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build(
                    "STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d INT>"
                ),
            },
            True,
        ),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build(
                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d INT>>"
                ),
            },
            True,
        ),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build(
                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d UNKNOWN>>"
                ),
            },
            False,
        ),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build(
                    "ARRAY<STRUCT<sub_a INT, sub_b INT, sub_c INT, sub_d UNKNOWN>>"
                ),
            },
            False,
        ),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build("MAP<INT, STRING>"),
            },
            True,
        ),
        (
            {
                "a": expressions.DataType.build("INT"),
                "b": expressions.DataType.build("MAP<INT, UNKNOWN>"),
            },
            False,
        ),
        (
            {"a": annotate_types(expressions.DataType.build("VARCHAR(MAX)", dialect="redshift"))},
            True,
        ),
    ],
)
def test_columns_to_types_all_known(columns_to_types, expected) -> None:
    assert columns_to_types_all_known(columns_to_types) == expected


def use_terminal_console(func):
    @wraps(func)
    def test_wrapper(*args, **kwargs):
        orig_console = get_console()
        try:
            new_console = TerminalConsole()
            new_console.console.width = 80
            new_console.console.no_color = True
            set_console(new_console)
            func(*args, **kwargs)
        finally:
            set_console(orig_console)

    return test_wrapper