Coverage for tests / test_enums.py: 100%

1import importlib 8 ctx1abcdefgh

2 

3import pytest 8 ctx1abcdefgh

4from sqlalchemy import create_mock_engine 8 ctx1abcdefgh

5from sqlalchemy.sql.type_api import TypeEngine 8 ctx1abcdefgh

6from sqlmodel import SQLModel 8 ctx1abcdefgh

7 

8from . import test_enums_models 8 ctx1abcdefgh

9 

10""" 8 ctx1abcdefgh

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 

19def pg_dump(sql: TypeEngine, *args, **kwargs): 8 ctx1abcdefgh

20 dialect = sql.compile(dialect=postgres_engine.dialect) 8 ctx1ijklmnop

21 sql_str = str(dialect).rstrip() 8 ctx1ijklmnop

22 if sql_str: 8 ctx1ijklmnop

23 print(sql_str + ";") 8 ctx1ijklmnop

24 

25 

26def sqlite_dump(sql: TypeEngine, *args, **kwargs): 8 ctx1abcdefgh

27 dialect = sql.compile(dialect=sqlite_engine.dialect) 8 ctx1qrstuvwx

28 sql_str = str(dialect).rstrip() 8 ctx1qrstuvwx

29 if sql_str: 8 ctx1qrstuvwx

30 print(sql_str + ";") 8 ctx1qrstuvwx

31 

32 

33postgres_engine = create_mock_engine("postgresql://", pg_dump) 8 ctx1abcdefgh

34sqlite_engine = create_mock_engine("sqlite://", sqlite_dump) 8 ctx1abcdefgh

35 

36 

37def test_postgres_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 8 ctx1abcdefgh

38 assert test_enums_models, "Ensure the models are imported and registered" 8 ctx1ijklmnop

39 importlib.reload(test_enums_models) 8 ctx1ijklmnop

40 SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False) 8 ctx1ijklmnop

41 

42 captured = capsys.readouterr() 8 ctx1ijklmnop

43 assert "CREATE TYPE myenum1 AS ENUM ('A', 'B');" in captured.out 8 ctx1ijklmnop

44 assert "CREATE TYPE myenum2 AS ENUM ('C', 'D');" in captured.out 8 ctx1ijklmnop

45 

46 

47def test_sqlite_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 8 ctx1abcdefgh

48 assert test_enums_models, "Ensure the models are imported and registered" 8 ctx1qrstuvwx

49 importlib.reload(test_enums_models) 8 ctx1qrstuvwx

50 SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False) 8 ctx1qrstuvwx

51 

52 captured = capsys.readouterr() 8 ctx1qrstuvwx

53 assert "enum_field VARCHAR(1) NOT NULL" in captured.out, captured 8 ctx1qrstuvwx

54 assert "CREATE TYPE" not in captured.out 8 ctx1qrstuvwx

55 

56 

57def test_json_schema_flat_model_pydantic_v2(): 8 ctx1abcdefgh

58 assert test_enums_models.FlatModel.model_json_schema() == { 8 ctx1yzABCDEF

59 "title": "FlatModel", 

60 "type": "object", 

61 "properties": { 

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

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

64 }, 

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

66 "$defs": { 

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

68 }, 

69 } 

70 

71 

72def test_json_schema_inherit_model_pydantic_v2(): 8 ctx1abcdefgh

73 assert test_enums_models.InheritModel.model_json_schema() == { 8 ctx1GHIJKLMN

74 "title": "InheritModel", 

75 "type": "object", 

76 "properties": { 

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

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

79 }, 

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

81 "$defs": { 

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

83 }, 

84 }