Coverage for tests/test_update_execution.py: 100%

1"""Update query construction, compilation, and execution acceptance tests.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import cast 

7 

8from snektest import assert_eq, assert_is_none, assert_raises, test 

9 

10from snekql import ( 

11 MISSING, 

12 Database, 

13 Integer, 

14 Model, 

15 Pending, 

16 QueryConstructionError, 

17 Text, 

18 insert, 

19 select, 

20 update, 

21) 

22from snekql.query import compile_write_sql 

23from tests.logging_helpers import NULL_LOGGER 

24 

25 

26@test(mark="fast") 

27def update_compilation_accepts_multiple_quoted_assignments() -> None: 

28 """Update compiles multiple assignments in caller-provided order.""" 

29 

30 class Order[S = Pending](Model[S, "Order[object]"]): 

31 """Table model with identifiers requiring SQLite quoting.""" 

32 

33 __tablename__ = "select" 

34 status: Order.Col[str] = Text(nullable=False) 

35 where: Order.Col[str] = Text(nullable=False) 

36 

37 query = ( 

38 update(Order) 

39 .set( 

40 Order.where.to("new"), 

41 Order.status.to("done"), 

42 ) 

43 .where(Order.where.ne("old")) 

44 ) 

45 

46 sql, params = compile_write_sql(query) 

47 

48 expected_sql = 'UPDATE "select" SET "where" = ?, "status" = ? WHERE ("where" != ?)' 

49 assert_eq(sql, expected_sql) 

50 assert_eq(params, ("new", "done", "old")) 

51 

52 

53@test(mark="fast") 

54def update_assignments_must_belong_to_target_model() -> None: 

55 """Update set() rejects assignments built from another table model.""" 

56 

57 class User[S = Pending](Model[S, "User[object]"]): 

58 """Target table model for ownership checks.""" 

59 

60 email: User.Col[str] = Text(nullable=False) 

61 

62 class AuditLog[S = Pending](Model[S, "AuditLog[object]"]): 

63 """Unrelated table model for ownership checks.""" 

64 

65 message: AuditLog.Col[str] = Text(nullable=False) 

66 

67 set_assignments = cast("Callable[..., object]", update(User).set) 

68 

69 with assert_raises(QueryConstructionError): 

70 _ = set_assignments(AuditLog.message.to("wrong table")) 

71 

72 

73@test(mark="fast") 

74def update_rejects_generated_and_primary_key_assignments() -> None: 

75 """Generated and primary key columns are not update-assignable.""" 

76 

77 class User[S = Pending](Model[S, "User[object]"]): 

78 """Table model with both protected and updateable columns.""" 

79 

80 account_id: User.Col[int] = Integer(primary_key=True) 

81 email: User.Col[str] = Text(nullable=False) 

82 revision: User.GenCol[int] = Integer(default=MISSING) 

83 

84 with assert_raises(QueryConstructionError): 

85 _ = update(User).set(User.revision.to(2)) 

86 

87 with assert_raises(QueryConstructionError): 

88 _ = update(User).set(User.account_id.to(2)) 

89 

90 

91@test(mark="medium") 

92async def update_execute_returns_none_and_supports_explicit_all() -> None: 

93 """tx.execute(update(...).all()) returns None after full-table update.""" 

94 

95 class User[S = Pending](Model[S, "User[object]"]): 

96 """Table model updated through the async runtime.""" 

97 

98 id: User.GenCol[int] = Integer(primary_key=True, default=MISSING) 

99 email: User.Col[str] = Text(nullable=False) 

100 status: User.Col[str] = Text(nullable=False, default="active") 

101 

102 database = await Database.initialize( 

103 NULL_LOGGER, database=":memory:", models=[User] 

104 ) 

105 try: 

106 async with database.transaction() as transaction: 

107 await transaction.execute(insert(User(email="a@example.com"))) 

108 await transaction.execute(insert(User(email="b@example.com"))) 

109 

110 result = await transaction.execute( 

111 update(User).set(User.status.to("disabled")).all(), 

112 ) 

113 statuses = await transaction.fetch_all( 

114 select(User.status).all().order_by(User.id.asc()), 

115 ) 

116 finally: 

117 await database.close() 

118 

119 assert_is_none(result) 

120 assert_eq(statuses, ["disabled", "disabled"])