Coverage for tests / test_main.py: 96%

1from typing import Annotated 8 ctx1456789!#

2 

3import pytest 8 ctx1456789!#

4from sqlalchemy.exc import IntegrityError 8 ctx1456789!#

5from sqlalchemy.orm import RelationshipProperty 8 ctx1456789!#

6from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select 8 ctx1456789!#

7 

8 

9def test_should_allow_duplicate_row_if_unique_constraint_is_not_passed(clear_sqlmodel): 8 ctx1456789!#

10 class Hero(SQLModel, table=True): 8 ctx1abcdefgh

11 id: int | None = Field(default=None, primary_key=True) 8 ctx1abcdefgh

12 name: str 8 ctx1abcdefgh

13 secret_name: str 8 ctx1abcdefgh

14 age: int | None = None 8 ctx1abcdefgh

15 

16 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1abcdefgh

17 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1abcdefgh

18 

19 engine = create_engine("sqlite://") 8 ctx1abcdefgh

20 

21 SQLModel.metadata.create_all(engine) 8 ctx1abcdefgh

22 

23 with Session(engine) as session: 8 ctx1abcdefgh

24 session.add(hero_1) 8 ctx1abcdefgh

25 session.commit() 8 ctx1abcdefgh

26 session.refresh(hero_1) 8 ctx1abcdefgh

27 

28 with Session(engine) as session: 8 ctx1abcdefgh

29 session.add(hero_2) 8 ctx1abcdefgh

30 session.commit() 8 ctx1abcdefgh

31 session.refresh(hero_2) 8 ctx1abcdefgh

32 

33 with Session(engine) as session: 8 ctx1abcdefgh

34 heroes = session.exec(select(Hero)).all() 8 ctx1abcdefgh

35 assert len(heroes) == 2 8 ctx1abcdefgh

36 assert heroes[0].name == heroes[1].name 8 ctx1abcdefgh

37 

38 

39def test_should_allow_duplicate_row_if_unique_constraint_is_false(clear_sqlmodel): 8 ctx1456789!#

40 class Hero(SQLModel, table=True): 8 ctx1ijklmnop

41 id: int | None = Field(default=None, primary_key=True) 8 ctx1ijklmnop

42 name: str 8 ctx1ijklmnop

43 secret_name: str = Field(unique=False) 8 ctx1ijklmnop

44 age: int | None = None 8 ctx1ijklmnop

45 

46 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1ijklmnop

47 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1ijklmnop

48 

49 engine = create_engine("sqlite://") 8 ctx1ijklmnop

50 

51 SQLModel.metadata.create_all(engine) 8 ctx1ijklmnop

52 

53 with Session(engine) as session: 8 ctx1ijklmnop

54 session.add(hero_1) 8 ctx1ijklmnop

55 session.commit() 8 ctx1ijklmnop

56 session.refresh(hero_1) 8 ctx1ijklmnop

57 

58 with Session(engine) as session: 8 ctx1ijklmnop

59 session.add(hero_2) 8 ctx1ijklmnop

60 session.commit() 8 ctx1ijklmnop

61 session.refresh(hero_2) 8 ctx1ijklmnop

62 

63 with Session(engine) as session: 8 ctx1ijklmnop

64 heroes = session.exec(select(Hero)).all() 8 ctx1ijklmnop

65 assert len(heroes) == 2 8 ctx1ijklmnop

66 assert heroes[0].name == heroes[1].name 8 ctx1ijklmnop

67 

68 

69def test_should_raise_exception_when_try_to_duplicate_row_if_unique_constraint_is_true( 8 ctx1456789!#

70 clear_sqlmodel, 

71): 

72 class Hero(SQLModel, table=True): 8 ctx1OPQRSTUV

73 id: int | None = Field(default=None, primary_key=True) 8 ctx1OPQRSTUV

74 name: str 8 ctx1OPQRSTUV

75 secret_name: str = Field(unique=True) 8 ctx1OPQRSTUV

76 age: int | None = None 8 ctx1OPQRSTUV

77 

78 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1OPQRSTUV

79 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 8 ctx1OPQRSTUV

80 

81 engine = create_engine("sqlite://") 8 ctx1OPQRSTUV

82 

83 SQLModel.metadata.create_all(engine) 8 ctx1OPQRSTUV

84 

85 with Session(engine) as session: 8 ctx1OPQRSTUV

86 session.add(hero_1) 8 ctx1OPQRSTUV

87 session.commit() 8 ctx1OPQRSTUV

88 session.refresh(hero_1) 8 ctx1OPQRSTUV

89 

90 with pytest.raises(IntegrityError): 8 ctx1OPQRSTUV

91 with Session(engine) as session: 8 ctx1OPQRSTUV

92 session.add(hero_2) 8 ctx1OPQRSTUV

93 session.commit() 8 ctx1OPQRSTUV

94 

95 

96def test_sa_relationship_property(clear_sqlmodel): 8 ctx1456789!#

97 """Test https://github.com/tiangolo/sqlmodel/issues/315#issuecomment-1272122306""" 

98 

99 class Team(SQLModel, table=True): 8 ctx1yzABCDEF

100 id: int | None = Field(default=None, primary_key=True) 8 ctx1yzABCDEF

101 name: str = Field(unique=True) 8 ctx1yzABCDEF

102 heroes: list["Hero"] = Relationship( # noqa: F821 8 ctx1yzABCDEF

103 sa_relationship=RelationshipProperty("Hero", back_populates="team") 

104 ) 

105 

106 class Hero(SQLModel, table=True): 8 ctx1yzABCDEF

107 id: int | None = Field(default=None, primary_key=True) 8 ctx1yzABCDEF

108 name: str = Field(unique=True) 8 ctx1yzABCDEF

109 team_id: int | None = Field(default=None, foreign_key="team.id") 8 ctx1yzABCDEF

110 team: Team | None = Relationship( 8 ctx1yzABCDEF

111 sa_relationship=RelationshipProperty("Team", back_populates="heroes") 

112 ) 

113 

114 team_preventers = Team(name="Preventers") 8 ctx1yzABCDEF

115 hero_rusty_man = Hero(name="Rusty-Man", team=team_preventers) 8 ctx1yzABCDEF

116 

117 engine = create_engine("sqlite://", echo=True) 8 ctx1yzABCDEF

118 

119 SQLModel.metadata.create_all(engine) 8 ctx1yzABCDEF

120 

121 with Session(engine) as session: 8 ctx1yzABCDEF

122 session.add(hero_rusty_man) 8 ctx1yzABCDEF

123 session.commit() 8 ctx1yzABCDEF

124 session.refresh(hero_rusty_man) 8 ctx1yzABCDEF

125 # The next statement should not raise an AttributeError 

126 assert hero_rusty_man.team 8 ctx1yzABCDEF

127 assert hero_rusty_man.team.name == "Preventers" 8 ctx1yzABCDEF

128 

129 

130def test_composite_primary_key(clear_sqlmodel): 8 ctx1456789!#

131 class UserPermission(SQLModel, table=True): 8 ctx1GHIJKLMN

132 user_id: int = Field(primary_key=True) 8 ctx1GHIJKLMN

133 resource_id: int = Field(primary_key=True) 8 ctx1GHIJKLMN

134 permission: str 8 ctx1GHIJKLMN

135 

136 engine = create_engine("sqlite://") 8 ctx1GHIJKLMN

137 SQLModel.metadata.create_all(engine) 8 ctx1GHIJKLMN

138 

139 pk_column_names = {column.name for column in UserPermission.__table__.primary_key} 8 ctx1GHIJKLMN

140 assert pk_column_names == {"user_id", "resource_id"} 8 ctx1GHIJKLMN

141 

142 with Session(engine) as session: 8 ctx1GHIJKLMN

143 perm1 = UserPermission(user_id=1, resource_id=1, permission="read") 8 ctx1GHIJKLMN

144 perm2 = UserPermission(user_id=1, resource_id=2, permission="write") 8 ctx1GHIJKLMN

145 session.add(perm1) 8 ctx1GHIJKLMN

146 session.add(perm2) 8 ctx1GHIJKLMN

147 session.commit() 8 ctx1GHIJKLMN

148 

149 with pytest.raises(IntegrityError): 8 ctx1GHIJKLMN

150 with Session(engine) as session: 8 ctx1GHIJKLMN

151 perm3 = UserPermission(user_id=1, resource_id=1, permission="admin") 8 ctx1GHIJKLMN

152 session.add(perm3) 8 ctx1GHIJKLMN

153 session.commit() 8 ctx1GHIJKLMN

154 

155 

156def test_composite_primary_key_and_validator(clear_sqlmodel): 8 ctx1456789!#

157 from pydantic import AfterValidator 8 ctx1qrstuvwx

158 

159 def validate_resource_id(value: int) -> int: 8 ctx1qrstuvwx

160 if value < 1: 

161 raise ValueError("Resource ID must be positive") 

162 return value 

163 

164 class UserPermission(SQLModel, table=True): 8 ctx1qrstuvwx

165 user_id: int = Field(primary_key=True) 8 ctx1qrstuvwx

166 resource_id: Annotated[int, AfterValidator(validate_resource_id)] = Field( 8 ctx1qrstuvwx

167 primary_key=True 

168 ) 

169 permission: str 8 ctx1qrstuvwx

170 

171 engine = create_engine("sqlite://") 8 ctx1qrstuvwx

172 SQLModel.metadata.create_all(engine) 8 ctx1qrstuvwx

173 

174 pk_column_names = {column.name for column in UserPermission.__table__.primary_key} 8 ctx1qrstuvwx

175 assert pk_column_names == {"user_id", "resource_id"} 8 ctx1qrstuvwx

176 

177 with Session(engine) as session: 8 ctx1qrstuvwx

178 perm1 = UserPermission(user_id=1, resource_id=1, permission="read") 8 ctx1qrstuvwx

179 perm2 = UserPermission(user_id=1, resource_id=2, permission="write") 8 ctx1qrstuvwx

180 session.add(perm1) 8 ctx1qrstuvwx

181 session.add(perm2) 8 ctx1qrstuvwx

182 session.commit() 8 ctx1qrstuvwx

183 

184 with pytest.raises(IntegrityError): 8 ctx1qrstuvwx

185 with Session(engine) as session: 8 ctx1qrstuvwx

186 perm3 = UserPermission(user_id=1, resource_id=1, permission="admin") 8 ctx1qrstuvwx

187 session.add(perm3) 8 ctx1qrstuvwx

188 session.commit() 8 ctx1qrstuvwx

189 

190 

191def test_foreign_key_ondelete_with_annotated(clear_sqlmodel): 8 ctx1456789!#

192 from pydantic import AfterValidator 8 ctx1WXYZ0123

193 

194 def ensure_positive(value: int) -> int: 8 ctx1WXYZ0123

195 if value < 0: 

196 raise ValueError("Team ID must be positive") 

197 return value 

198 

199 class Team(SQLModel, table=True): 8 ctx1WXYZ0123

200 id: int = Field(primary_key=True) 8 ctx1WXYZ0123

201 name: str 8 ctx1WXYZ0123

202 

203 class Hero(SQLModel, table=True): 8 ctx1WXYZ0123

204 id: int = Field(primary_key=True) 8 ctx1WXYZ0123

205 team_id: Annotated[int, AfterValidator(ensure_positive)] = Field( 8 ctx1WXYZ0123

206 foreign_key="team.id", 

207 ondelete="CASCADE", 

208 ) 

209 name: str 8 ctx1WXYZ0123

210 

211 engine = create_engine("sqlite://") 8 ctx1WXYZ0123

212 SQLModel.metadata.create_all(engine) 8 ctx1WXYZ0123

213 

214 team_id_column = Hero.__table__.c.team_id # type: ignore[attr-defined] 8 ctx1WXYZ0123

215 foreign_keys = list(team_id_column.foreign_keys) 8 ctx1WXYZ0123

216 assert len(foreign_keys) == 1 8 ctx1WXYZ0123

217 assert foreign_keys[0].ondelete == "CASCADE" 8 ctx1WXYZ0123

218 assert team_id_column.nullable is False 8 ctx1WXYZ0123