Coverage for tests/test_enums.py: 100%

1import enum 6 ctx1fabcde

2import uuid 6 ctx1fabcde

3 

4from sqlalchemy import create_mock_engine 6 ctx1fabcde

5from sqlalchemy.sql.type_api import TypeEngine 6 ctx1fabcde

6from sqlmodel import Field, SQLModel 6 ctx1fabcde

7 

8from .conftest import needs_pydanticv1, needs_pydanticv2 6 ctx1fabcde

9 

10""" 

11Tests related to Enums 

12 

13Associated issues: 

14* https://github.com/tiangolo/sqlmodel/issues/96 

15* https://github.com/tiangolo/sqlmodel/issues/164 

16""" 

17 

18 

19class MyEnum1(str, enum.Enum): 6 ctx1fabcde

20 A = "A" 6 ctx1fabcde

21 B = "B" 6 ctx1fabcde

22 

23 

24class MyEnum2(str, enum.Enum): 6 ctx1fabcde

25 C = "C" 6 ctx1fabcde

26 D = "D" 6 ctx1fabcde

27 

28 

29class BaseModel(SQLModel): 6 ctx1fabcde

30 id: uuid.UUID = Field(primary_key=True) 6 ctx1fabcde

31 enum_field: MyEnum2 6 ctx1fabcde

32 

33 

34class FlatModel(SQLModel, table=True): 6 ctx1fabcde

35 id: uuid.UUID = Field(primary_key=True) 6 ctx1fabcde

36 enum_field: MyEnum1 6 ctx1fabcde

37 

38 

39class InheritModel(BaseModel, table=True): 6 ctx1fabcde

40 pass 6 ctx1fabcde

41 

42 

43def pg_dump(sql: TypeEngine, *args, **kwargs): 6 ctx1fabcde

44 dialect = sql.compile(dialect=postgres_engine.dialect) 6 ctx1fabcde

45 sql_str = str(dialect).rstrip() 6 ctx1fabcde

46 if sql_str: 6 ctx1fabcde

47 print(sql_str + ";") 6 ctx1fabcde

48 

49 

50def sqlite_dump(sql: TypeEngine, *args, **kwargs): 6 ctx1fabcde

51 dialect = sql.compile(dialect=sqlite_engine.dialect) 6 ctx1fabcde

52 sql_str = str(dialect).rstrip() 6 ctx1fabcde

53 if sql_str: 6 ctx1fabcde

54 print(sql_str + ";") 6 ctx1fabcde

55 

56 

57postgres_engine = create_mock_engine("postgresql://", pg_dump) 6 ctx1fabcde

58sqlite_engine = create_mock_engine("sqlite://", sqlite_dump) 6 ctx1fabcde

59 

60 

61def test_postgres_ddl_sql(capsys): 6 ctx1fabcde

62 SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False) 6 ctx1fabcde

63 

64 captured = capsys.readouterr() 6 ctx1fabcde

65 assert "CREATE TYPE myenum1 AS ENUM ('A', 'B');" in captured.out 6 ctx1fabcde

66 assert "CREATE TYPE myenum2 AS ENUM ('C', 'D');" in captured.out 6 ctx1fabcde

67 

68 

69def test_sqlite_ddl_sql(capsys): 6 ctx1fabcde

70 SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False) 6 ctx1fabcde

71 

72 captured = capsys.readouterr() 6 ctx1fabcde

73 assert "enum_field VARCHAR(1) NOT NULL" in captured.out 6 ctx1fabcde

74 assert "CREATE TYPE" not in captured.out 6 ctx1fabcde

75 

76 

77@needs_pydanticv1 6 ctx1fabcde

78def test_json_schema_flat_model_pydantic_v1(): 5 ctx1abcde

79 assert FlatModel.schema() == { 4 ctx1bcde

80 "title": "FlatModel", 

81 "type": "object", 

82 "properties": { 

83 "id": {"title": "Id", "type": "string", "format": "uuid"}, 

84 "enum_field": {"$ref": "#/definitions/MyEnum1"}, 

85 }, 

86 "required": ["id", "enum_field"], 

87 "definitions": { 

88 "MyEnum1": { 

89 "title": "MyEnum1", 

90 "description": "An enumeration.", 

91 "enum": ["A", "B"], 

92 "type": "string", 

93 } 

94 }, 

95 } 

96 

97 

98@needs_pydanticv1 6 ctx1fabcde

99def test_json_schema_inherit_model_pydantic_v1(): 5 ctx1abcde

100 assert InheritModel.schema() == { 4 ctx1bcde

101 "title": "InheritModel", 

102 "type": "object", 

103 "properties": { 

104 "id": {"title": "Id", "type": "string", "format": "uuid"}, 

105 "enum_field": {"$ref": "#/definitions/MyEnum2"}, 

106 }, 

107 "required": ["id", "enum_field"], 

108 "definitions": { 

109 "MyEnum2": { 

110 "title": "MyEnum2", 

111 "description": "An enumeration.", 

112 "enum": ["C", "D"], 

113 "type": "string", 

114 } 

115 }, 

116 } 

117 

118 

119@needs_pydanticv2 6 ctx1fabcde

120def test_json_schema_flat_model_pydantic_v2(): 5 ctx1abcde

121 assert FlatModel.model_json_schema() == { 2 ctx1fa

122 "title": "FlatModel", 

123 "type": "object", 

124 "properties": { 

125 "id": {"title": "Id", "type": "string", "format": "uuid"}, 

126 "enum_field": {"$ref": "#/$defs/MyEnum1"}, 

127 }, 

128 "required": ["id", "enum_field"], 

129 "$defs": { 

130 "MyEnum1": {"enum": ["A", "B"], "title": "MyEnum1", "type": "string"} 

131 }, 

132 } 

133 

134 

135@needs_pydanticv2 6 ctx1fabcde

136def test_json_schema_inherit_model_pydantic_v2(): 5 ctx1abcde

137 assert InheritModel.model_json_schema() == { 2 ctx1fa

138 "title": "InheritModel", 

139 "type": "object", 

140 "properties": { 

141 "id": {"title": "Id", "type": "string", "format": "uuid"}, 

142 "enum_field": {"$ref": "#/$defs/MyEnum2"}, 

143 }, 

144 "required": ["id", "enum_field"], 

145 "$defs": { 

146 "MyEnum2": {"enum": ["C", "D"], "title": "MyEnum2", "type": "string"} 

147 }, 

148 }