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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150
|
"""Tests for Pydantic v1 BaseModel generation."""
from __future__ import annotations
from typing import Any
import pytest
from datamodel_code_generator.model.pydantic import DataTypeManager
from datamodel_code_generator.model.pydantic.base_model import BaseModel, DataModelField
from datamodel_code_generator.reference import Reference
from datamodel_code_generator.types import DataType, Types
def test_base_model() -> None:
"""Test basic BaseModel generation with required field."""
field = DataModelField(name="a", data_type=DataType(type="str"), required=True)
base_model = BaseModel(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert base_model.name == "test_model"
assert base_model.fields == [field]
assert base_model.decorators == []
assert base_model.render() == "class test_model(BaseModel):\n a: str"
def test_base_model_optional() -> None:
"""Test BaseModel generation with optional field and default value."""
field = DataModelField(name="a", data_type=DataType(type="str"), default="abc", required=False)
base_model = BaseModel(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert base_model.name == "test_model"
assert base_model.fields == [field]
assert base_model.decorators == []
assert base_model.render() == "class test_model(BaseModel):\n a: Optional[str] = 'abc'"
def test_base_model_nullable_required() -> None:
"""Test BaseModel generation with nullable required field."""
field = DataModelField(
name="a",
data_type=DataType(type="str"),
default="abc",
required=True,
nullable=True,
)
base_model = BaseModel(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert base_model.name == "test_model"
assert base_model.fields == [field]
assert base_model.decorators == []
assert base_model.render() == "class test_model(BaseModel):\n a: Optional[str] = Field(...)"
def test_base_model_strict_non_nullable_required() -> None:
"""Test BaseModel generation with strict non-nullable required field."""
field = DataModelField(
name="a",
data_type=DataType(type="str"),
default="abc",
required=True,
nullable=False,
)
base_model = BaseModel(
fields=[field],
reference=Reference(name="test_model", path="test_model"),
)
assert base_model.name == "test_model"
assert base_model.fields == [field]
assert base_model.decorators == []
assert base_model.render() == "class test_model(BaseModel):\n a: str"
def test_base_model_decorator() -> None:
"""Test BaseModel generation with decorators and base classes."""
field = DataModelField(name="a", data_type=DataType(type="str"), default="abc", required=False)
base_model = BaseModel(
fields=[field],
decorators=["@validate"],
base_classes=[Reference(name="Base", original_name="Base", path="Base")],
reference=Reference(name="test_model", path="test_model"),
)
assert base_model.name == "test_model"
assert base_model.fields == [field]
assert base_model.base_class == "Base"
assert base_model.decorators == ["@validate"]
assert base_model.render() == "@validate\nclass test_model(Base):\n a: Optional[str] = 'abc'"
def test_base_model_get_data_type() -> None:
"""Test data type retrieval for BaseModel fields."""
data_type_manager = DataTypeManager()
assert data_type_manager.get_data_type(Types.integer) == data_type_manager.data_type(type="int")
@pytest.mark.parametrize(
("kwargs", "expected"),
[
({"required": True}, None),
({"required": True, "example": "example"}, None),
({"example": "example"}, None),
(
{"required": True, "default": 123, "example": "example"},
None,
),
(
{"required": False, "default": 123, "example": "example"},
None,
),
({"description": "description"}, None),
({"title": "title"}, None),
({"alias": "alias"}, "Field(None, alias='alias')"),
({"example": True}, None),
({"examples": True}, None),
(
{
"example": True,
"description": "description",
"title": "title",
"alias": "alias",
},
"Field(None, alias='alias')",
),
({"examples": [1, 2, 3]}, None),
(
{"examples": {"name": "dog", "age": 1}},
None,
),
({"default": "abc", "title": "title"}, None),
({"default": 123, "title": "title"}, None),
],
)
def test_data_model_field(kwargs: dict[str, Any], expected: str | None) -> None:
"""Test DataModelField generation with various field attributes."""
assert DataModelField(**kwargs, data_type=DataType()).field == expected
|