Coverage for tests / test_main.py: 96%

153 statements  

« prev     ^ index     » next       coverage.py v7.13.0, created at 2026-11-09 00:02 +0000

1from typing import Annotated 1#$%'()*+,

2 

3import pytest 1#$%'()*+,

4from sqlalchemy import Engine 1#$%'()*+,

5from sqlalchemy.exc import IntegrityError 1#$%'()*+,

6from sqlalchemy.orm import RelationshipProperty 1#$%'()*+,

7from sqlmodel import Field, Relationship, Session, SQLModel, select 1#$%'()*+,

8 

9 

10def test_should_allow_duplicate_row_if_unique_constraint_is_not_passed( 1#$%'()*+,

11 database_engine: Engine, 

12): 

13 class Hero(SQLModel, table=True): 1abcdefghi

14 id: int | None = Field(default=None, primary_key=True) 1abcdefghi

15 name: str 1abcdefghi

16 secret_name: str 1abcdefghi

17 age: int | None = None 1abcdefghi

18 

19 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 1abcdefghi

20 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 1abcdefghi

21 

22 engine = database_engine 1abcdefghi

23 

24 SQLModel.metadata.create_all(engine) 1abcdefghi

25 

26 with Session(engine) as session: 1abcdefghi

27 session.add(hero_1) 1abcdefghi

28 session.commit() 1abcdefghi

29 session.refresh(hero_1) 1abcdefghi

30 

31 with Session(engine) as session: 1abcdefghi

32 session.add(hero_2) 1abcdefghi

33 session.commit() 1abcdefghi

34 session.refresh(hero_2) 1abcdefghi

35 

36 with Session(engine) as session: 1abcdefghi

37 heroes = session.exec(select(Hero)).all() 1abcdefghi

38 assert len(heroes) == 2 1abcdefghi

39 assert heroes[0].name == heroes[1].name 1abcdefghi

40 

41 

42def test_should_allow_duplicate_row_if_unique_constraint_is_false( 1#$%'()*+,

43 database_engine: Engine, 

44): 

45 class Hero(SQLModel, table=True): 1jklmnopqr

46 id: int | None = Field(default=None, primary_key=True) 1jklmnopqr

47 name: str 1jklmnopqr

48 secret_name: str = Field(unique=False) 1jklmnopqr

49 age: int | None = None 1jklmnopqr

50 

51 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 1jklmnopqr

52 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 1jklmnopqr

53 

54 engine = database_engine 1jklmnopqr

55 

56 SQLModel.metadata.create_all(engine) 1jklmnopqr

57 

58 with Session(engine) as session: 1jklmnopqr

59 session.add(hero_1) 1jklmnopqr

60 session.commit() 1jklmnopqr

61 session.refresh(hero_1) 1jklmnopqr

62 

63 with Session(engine) as session: 1jklmnopqr

64 session.add(hero_2) 1jklmnopqr

65 session.commit() 1jklmnopqr

66 session.refresh(hero_2) 1jklmnopqr

67 

68 with Session(engine) as session: 1jklmnopqr

69 heroes = session.exec(select(Hero)).all() 1jklmnopqr

70 assert len(heroes) == 2 1jklmnopqr

71 assert heroes[0].name == heroes[1].name 1jklmnopqr

72 

73 

74def test_should_raise_exception_when_try_to_duplicate_row_if_unique_constraint_is_true( 1#$%'()*+,

75 database_engine: Engine, 

76): 

77 class Hero(SQLModel, table=True): 1TUVWXYZ01

78 id: int | None = Field(default=None, primary_key=True) 1TUVWXYZ01

79 name: str 1TUVWXYZ01

80 secret_name: str = Field(unique=True) 1TUVWXYZ01

81 age: int | None = None 1TUVWXYZ01

82 

83 hero_1 = Hero(name="Deadpond", secret_name="Dive Wilson") 1TUVWXYZ01

84 hero_2 = Hero(name="Deadpond", secret_name="Dive Wilson") 1TUVWXYZ01

85 

86 engine = database_engine 1TUVWXYZ01

87 

88 SQLModel.metadata.create_all(engine) 1TUVWXYZ01

89 

90 with Session(engine) as session: 1TUVWXYZ01

91 session.add(hero_1) 1TUVWXYZ01

92 session.commit() 1TUVWXYZ01

93 session.refresh(hero_1) 1TUVWXYZ01

94 

95 with pytest.raises(IntegrityError): 1TUVWXYZ01

96 with Session(engine) as session: 1TUVWXYZ01

97 session.add(hero_2) 1TUVWXYZ01

98 session.commit() 1TUVWXYZ01

99 

100 

101def test_sa_relationship_property(database_engine: Engine): 1#$%'()*+,

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

103 

104 class Team(SQLModel, table=True): 1BCDEFGHIJ

105 id: int | None = Field(default=None, primary_key=True) 1BCDEFGHIJ

106 name: str = Field(unique=True) 1BCDEFGHIJ

107 heroes: list["Hero"] = Relationship( # noqa: F821 1BCDEFGHIJ

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

109 ) 

110 

111 class Hero(SQLModel, table=True): 1BCDEFGHIJ

112 id: int | None = Field(default=None, primary_key=True) 1BCDEFGHIJ

113 name: str = Field(unique=True) 1BCDEFGHIJ

114 team_id: int | None = Field(default=None, foreign_key="team.id") 1BCDEFGHIJ

115 team: Team | None = Relationship( 1BCDEFGHIJ

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

117 ) 

118 

119 team_preventers = Team(name="Preventers") 1BCDEFGHIJ

120 hero_rusty_man = Hero(name="Rusty-Man", team=team_preventers) 1BCDEFGHIJ

121 

122 engine = database_engine 1BCDEFGHIJ

123 

124 SQLModel.metadata.create_all(engine) 1BCDEFGHIJ

125 

126 with Session(engine) as session: 1BCDEFGHIJ

127 session.add(hero_rusty_man) 1BCDEFGHIJ

128 session.commit() 1BCDEFGHIJ

129 session.refresh(hero_rusty_man) 1BCDEFGHIJ

130 # The next statement should not raise an AttributeError 

131 assert hero_rusty_man.team 1BCDEFGHIJ

132 assert hero_rusty_man.team.name == "Preventers" 1BCDEFGHIJ

133 

134 

135def test_composite_primary_key(database_engine: Engine): 1#$%'()*+,

136 class UserPermission(SQLModel, table=True): 1KLMNOPQRS

137 user_id: int = Field(primary_key=True) 1KLMNOPQRS

138 resource_id: int = Field(primary_key=True) 1KLMNOPQRS

139 permission: str 1KLMNOPQRS

140 

141 engine = database_engine 1KLMNOPQRS

142 SQLModel.metadata.create_all(engine) 1KLMNOPQRS

143 

144 pk_column_names = {column.name for column in UserPermission.__table__.primary_key} 1KLMNOPQRS

145 assert pk_column_names == {"user_id", "resource_id"} 1KLMNOPQRS

146 

147 with Session(engine) as session: 1KLMNOPQRS

148 perm1 = UserPermission(user_id=1, resource_id=1, permission="read") 1KLMNOPQRS

149 perm2 = UserPermission(user_id=1, resource_id=2, permission="write") 1KLMNOPQRS

150 session.add(perm1) 1KLMNOPQRS

151 session.add(perm2) 1KLMNOPQRS

152 session.commit() 1KLMNOPQRS

153 

154 with pytest.raises(IntegrityError): 1KLMNOPQRS

155 with Session(engine) as session: 1KLMNOPQRS

156 perm3 = UserPermission(user_id=1, resource_id=1, permission="admin") 1KLMNOPQRS

157 session.add(perm3) 1KLMNOPQRS

158 session.commit() 1KLMNOPQRS

159 

160 

161def test_composite_primary_key_and_validator(database_engine: Engine): 1#$%'()*+,

162 from pydantic import AfterValidator 1stuvwxyzA

163 

164 def validate_resource_id(value: int) -> int: 1stuvwxyzA

165 if value < 1: 

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

167 return value 

168 

169 class UserPermission(SQLModel, table=True): 1stuvwxyzA

170 user_id: int = Field(primary_key=True) 1stuvwxyzA

171 resource_id: Annotated[int, AfterValidator(validate_resource_id)] = Field( 1stuvwxyzA

172 primary_key=True 

173 ) 

174 permission: str 1stuvwxyzA

175 

176 engine = database_engine 1stuvwxyzA

177 SQLModel.metadata.create_all(engine) 1stuvwxyzA

178 

179 pk_column_names = {column.name for column in UserPermission.__table__.primary_key} 1stuvwxyzA

180 assert pk_column_names == {"user_id", "resource_id"} 1stuvwxyzA

181 

182 with Session(engine) as session: 1stuvwxyzA

183 perm1 = UserPermission(user_id=1, resource_id=1, permission="read") 1stuvwxyzA

184 perm2 = UserPermission(user_id=1, resource_id=2, permission="write") 1stuvwxyzA

185 session.add(perm1) 1stuvwxyzA

186 session.add(perm2) 1stuvwxyzA

187 session.commit() 1stuvwxyzA

188 

189 with pytest.raises(IntegrityError): 1stuvwxyzA

190 with Session(engine) as session: 1stuvwxyzA

191 perm3 = UserPermission(user_id=1, resource_id=1, permission="admin") 1stuvwxyzA

192 session.add(perm3) 1stuvwxyzA

193 session.commit() 1stuvwxyzA

194 

195 

196def test_foreign_key_ondelete_with_annotated(database_engine: Engine): 1#$%'()*+,

197 from pydantic import AfterValidator 123456789!

198 

199 def ensure_positive(value: int) -> int: 123456789!

200 if value < 0: 

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

202 return value 

203 

204 class Team(SQLModel, table=True): 123456789!

205 id: int = Field(primary_key=True) 123456789!

206 name: str 123456789!

207 

208 class Hero(SQLModel, table=True): 123456789!

209 id: int = Field(primary_key=True) 123456789!

210 team_id: Annotated[int, AfterValidator(ensure_positive)] = Field( 123456789!

211 foreign_key="team.id", 

212 ondelete="CASCADE", 

213 ) 

214 name: str 123456789!

215 

216 engine = database_engine 123456789!

217 SQLModel.metadata.create_all(engine) 123456789!

218 

219 team_id_column = Hero.__table__.c.team_id # type: ignore[attr-defined] 123456789!

220 foreign_keys = list(team_id_column.foreign_keys) 123456789!

221 assert len(foreign_keys) == 1 123456789!

222 assert foreign_keys[0].ondelete == "CASCADE" 123456789!

223 assert team_id_column.nullable is False 123456789!