Coverage for tests / test_enums.py: 100%

1import importlib 9 ctx1abcidefgh

2 

3import pytest 9 ctx1abcidefgh

4from sqlalchemy import create_mock_engine 9 ctx1abcidefgh

5from sqlalchemy.sql.type_api import TypeEngine 9 ctx1abcidefgh

6from sqlmodel import SQLModel 9 ctx1abcidefgh

7 

8from . import test_enums_models 9 ctx1abcidefgh

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): 9 ctx1abcidefgh

20 dialect = sql.compile(dialect=postgres_engine.dialect) 9 ctx1jklmnopqr

21 sql_str = str(dialect).rstrip() 9 ctx1jklmnopqr

22 if sql_str: 9 ctx1jklmnopqr

23 print(sql_str + ";") 9 ctx1jklmnopqr

24 

25 

26def sqlite_dump(sql: TypeEngine, *args, **kwargs): 9 ctx1abcidefgh

27 dialect = sql.compile(dialect=sqlite_engine.dialect) 9 ctx1stuvwxyzA

28 sql_str = str(dialect).rstrip() 9 ctx1stuvwxyzA

29 if sql_str: 9 ctx1stuvwxyzA

30 print(sql_str + ";") 9 ctx1stuvwxyzA

31 

32 

33postgres_engine = create_mock_engine("postgresql://", pg_dump) 9 ctx1abcidefgh

34sqlite_engine = create_mock_engine("sqlite://", sqlite_dump) 9 ctx1abcidefgh

35 

36 

37def test_postgres_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 9 ctx1abcidefgh

38 assert test_enums_models, "Ensure the models are imported and registered" 9 ctx1jklmnopqr

39 importlib.reload(test_enums_models) 9 ctx1jklmnopqr

40 SQLModel.metadata.create_all(bind=postgres_engine, checkfirst=False) 9 ctx1jklmnopqr

41 

42 captured = capsys.readouterr() 9 ctx1jklmnopqr

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

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

45 

46 

47def test_sqlite_ddl_sql(clear_sqlmodel, capsys: pytest.CaptureFixture[str]): 9 ctx1abcidefgh

48 assert test_enums_models, "Ensure the models are imported and registered" 9 ctx1stuvwxyzA

49 importlib.reload(test_enums_models) 9 ctx1stuvwxyzA

50 SQLModel.metadata.create_all(bind=sqlite_engine, checkfirst=False) 9 ctx1stuvwxyzA

51 

52 captured = capsys.readouterr() 9 ctx1stuvwxyzA

53 assert "enum_field VARCHAR(1) NOT NULL" in captured.out, captured 9 ctx1stuvwxyzA

54 assert "CREATE TYPE" not in captured.out 9 ctx1stuvwxyzA

55 

56 

57def test_json_schema_flat_model_pydantic_v2(): 9 ctx1abcidefgh

58 assert test_enums_models.FlatModel.model_json_schema() == { 9 ctx1BCDEFGHIJ

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(): 9 ctx1abcidefgh

73 assert test_enums_models.InheritModel.model_json_schema() == { 9 ctx1KLMNOPQRS

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 }