cli/python/codegen/codegen_tests/test_generated_dataclass.py at e9158db29e05b6e3f97dbb17c712f3e4ba6cd198 · databricks/cli · 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
import codegen.packages as packages
from codegen.generated_dataclass import (
    GeneratedDataclass,
    GeneratedField,
    dict_type,
    generate_dataclass,
    generate_type,
    GeneratedType,
    str_type,
    variable_or_type,
)
from codegen.jsonschema import Property, Schema, SchemaType


def test_generate_type_string():
    generated_type = generate_type(
        namespace="jobs",
        ref="#/$defs/string",
        is_param=False,
    )

    assert generated_type == GeneratedType(
        name="str",
        package=None,
        parameters=[],
    )


def test_generate_type_dict():
    generated_type = generate_type(
        namespace="jobs",
        ref="#/$defs/map/string",
        is_param=False,
    )

    assert generated_type == dict_type()


def test_generate_dataclass():
    generated = generate_dataclass(
        namespace="bananas",
        schema_name="jobs.Task",
        schema=Schema(
            type=SchemaType.OBJECT,
            description="task description",
            properties={
                "task_key": Property(
                    description="task_key description",
                    ref="#/$defs/string",
                ),
            },
            required=["task_key"],
        ),
    )

    assert generated == GeneratedDataclass(
        class_name="Task",
        package="databricks.bundles.bananas._models.task",
        description="task description",
        extends=[],
        fields=[
            GeneratedField(
                create_func_default=None,
                create_func_type_name=variable_or_type(str_type(), is_required=True),
                default=None,
                default_factory=None,
                description="task_key description",
                field_name="task_key",
                param_type_name=variable_or_type(str_type(), is_required=True),
                type_name=variable_or_type(str_type(), is_required=True),
                experimental=False,
                deprecated=False,
            ),
        ],
        experimental=False,
        deprecated=False,
    )


def test_get_class_name_string():
    class_name = packages.get_class_name("#/$defs/string")

    assert class_name == "str"