Coverage for tests/test_enums.py: 100%

1import importlib 6 ctx1fdeabc

2 

3import pytest 6 ctx1fdeabc

4from sqlalchemy import create_mock_engine 6 ctx1fdeabc

5from sqlalchemy.sql.type_api import TypeEngine 6 ctx1fdeabc

6from sqlmodel import SQLModel 6 ctx1fdeabc

7 

8from . import test_enums_models 6 ctx1fdeabc

9from .conftest import needs_pydanticv1, needs_pydanticv2 6 ctx1fdeabc

10 

11""" 3 ctx1abc

12Tests related to Enums 

13 

14Associated issues: 

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

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

17""" 

18 

19 

20def pg_dump(sql: TypeEngine, *args, **kwargs): 6 ctx1fdeabc

21 dialect = sql.compile(dialect=postgres_engine.dialect) 6 ctx1ghijkl

22 sql_str = str(dialect).rstrip() 6 ctx1ghijkl

23 if sql_str: 6 ctx1ghijkl

24 print(sql_str + ";") 6 ctx1ghijkl

25 

26 

27def sqlite_dump(sql: TypeEngine, *args, **kwargs): 6 ctx1fdeabc

28 dialect = sql.compile(dialect=sqlite_engine.dialect) 6 ctx1mnopqr

29 sql_str = str(dialect).rstrip() 6 ctx1mnopqr

30 if sql_str: 6 ctx1mnopqr

31 print(sql_str + ";") 6 ctx1mnopqr

32 

33 

34postgres_engine = create_mock_engine("postgresql://", pg_dump) 6 ctx1fdeabc

35sqlite_engine = create_mock_engine("sqlite://", sqlite_dump) 6 ctx1fdeabc

36 

37 

38def test_postgres_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 6 ctx1fdeabc

39 assert test_enums_models, "Ensure the models are imported and registered" 6 ctx1ghijkl

40 importlib.reload(test_enums_models) 6 ctx1ghijkl

41 SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False) 6 ctx1ghijkl

42 

43 captured = capsys.readouterr() 6 ctx1ghijkl

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

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

46 

47 

48def test_sqlite_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 6 ctx1fdeabc

49 assert test_enums_models, "Ensure the models are imported and registered" 6 ctx1mnopqr

50 importlib.reload(test_enums_models) 6 ctx1mnopqr

51 SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False) 6 ctx1mnopqr

52 

53 captured = capsys.readouterr() 6 ctx1mnopqr

54 assert "enum_field VARCHAR(1) NOT NULL" in captured.out, captured 6 ctx1mnopqr

55 assert "CREATE TYPE" not in captured.out 6 ctx1mnopqr

56 

57 

58@needs_pydanticv1 6 ctx1fdeabc

59def test_json_schema_flat_model_pydantic_v1(): 5 ctx1deabc

60 assert test_enums_models.FlatModel.schema() == { 6 ctx1stuvwx

61 "title": "FlatModel", 

62 "type": "object", 

63 "properties": { 

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

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

66 }, 

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

68 "definitions": { 

69 "MyEnum1": { 

70 "title": "MyEnum1", 

71 "description": "An enumeration.", 

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

73 "type": "string", 

74 } 

75 }, 

76 } 

77 

78 

79@needs_pydanticv1 6 ctx1fdeabc

80def test_json_schema_inherit_model_pydantic_v1(): 5 ctx1deabc

81 assert test_enums_models.InheritModel.schema() == { 6 ctx1yzABCD

82 "title": "InheritModel", 

83 "type": "object", 

84 "properties": { 

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

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

87 }, 

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

89 "definitions": { 

90 "MyEnum2": { 

91 "title": "MyEnum2", 

92 "description": "An enumeration.", 

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

94 "type": "string", 

95 } 

96 }, 

97 } 

98 

99 

100@needs_pydanticv2 6 ctx1fdeabc

101def test_json_schema_flat_model_pydantic_v2(): 5 ctx1deabc

102 assert test_enums_models.FlatModel.model_json_schema() == { 6 ctx1EFGHIJ

103 "title": "FlatModel", 

104 "type": "object", 

105 "properties": { 

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

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

108 }, 

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

110 "$defs": { 

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

112 }, 

113 } 

114 

115 

116@needs_pydanticv2 6 ctx1fdeabc

117def test_json_schema_inherit_model_pydantic_v2(): 5 ctx1deabc

118 assert test_enums_models.InheritModel.model_json_schema() == { 6 ctx1KLMNOP

119 "title": "InheritModel", 

120 "type": "object", 

121 "properties": { 

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

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

124 }, 

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

126 "$defs": { 

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

128 }, 

129 }