Coverage for app/backend/src/tests/fixtures/timewarp.py: 98%

1from collections.abc import Generator 

2from datetime import UTC, datetime, timedelta 

3from unittest.mock import patch 

4 

5from sqlalchemy import Connection, event, text 

6 

7from couchers import utils 

8from couchers.db import _get_base_engine 

9 

10FROZEN_TEST_TIME = datetime(2020, 1, 1, tzinfo=UTC) 

11 

12 

13class Clock: 

14 """ 

15 What the `timewarp` and `frozen_timewarp` fixtures below have in common: a clock that python and 

16 postgres both read, and that a test moves with advance(timedelta(days=30)), negative to rewind. 

17 

18 Whether it ticks is a property of the fixture the test asked for rather than something it can 

19 turn on and off partway through, so the two are separate clocks rather than one with a mode. 

20 Each puts the clock at an instant under its own name, run_from and freeze_at, so a call site 

21 says which kind of clock it's talking to. 

22 """ 

23 

24 def __init__(self) -> None: 

25 self._open_transactions: set[Connection] = set() 

26 

27 def now(self) -> datetime: 

28 raise NotImplementedError 

29 

30 def _db_settings(self) -> tuple[str, str]: 

31 """What postgres needs to read this same clock, as (mock.offset, mock.frozen_at).""" 

32 raise NotImplementedError 

33 

34 def _refuse_mid_transaction(self) -> None: 

35 if self._open_transactions: 

36 raise RuntimeError( 

37 "can't move the clock while a transaction is open: postgres reads the clock as it " 

38 "was at transaction start, so this would silently do nothing until the next one. " 

39 "Move the clock outside the session_scope block." 

40 ) 

41 

42 

43class Timewarp(Clock): 

44 """ 

45 A clock that keeps running, displaced from the real one by an offset. Each side applies that 

46 offset to its own clock, so the two stay as (im)perfectly aligned as they normally are, and 

47 run_from lands within a hair of the instant asked for rather than exactly on it. Postgres also 

48 still reports transaction start time, so a long transaction sees its own start rather than a 

49 later advance. 

50 """ 

51 

52 def __init__(self) -> None: 

53 super().__init__() 

54 self.offset = timedelta() 

55 

56 def now(self) -> datetime: 

57 return datetime.now(tz=UTC) + self.offset 

58 

59 def run_from(self, when: datetime) -> None: 

60 """Sets the clock to `when` and lets it tick on from there.""" 

61 _check_aware(when) 

62 self._refuse_mid_transaction() 

63 self.offset = when - datetime.now(tz=UTC) 

64 

65 def advance(self, delta: timedelta) -> None: 

66 """Moves the clock forwards, or backwards if delta is negative.""" 

67 self._refuse_mid_transaction() 

68 self.offset += delta 

69 

70 def _db_settings(self) -> tuple[str, str]: 

71 return _as_interval(self.offset), "" 

72 

73 

74class FrozenTimewarp(Clock): 

75 """ 

76 A clock stopped dead at one instant, so both sides read back exactly that rather than ticking on 

77 from there. freeze_at and advance move it to another standstill. 

78 """ 

79 

80 def __init__(self, at: datetime) -> None: 

81 super().__init__() 

82 _check_aware(at) 

83 # utils.now() is UTC-aware everywhere else, and callers do read the tzinfo off it 

84 self.frozen_at = at.astimezone(UTC) 

85 

86 def now(self) -> datetime: 

87 return self.frozen_at 

88 

89 def freeze_at(self, when: datetime) -> None: 

90 """Stops the clock at `when` instead of wherever it was stopped before.""" 

91 _check_aware(when) 

92 self._refuse_mid_transaction() 

93 self.frozen_at = when.astimezone(UTC) 

94 

95 def advance(self, delta: timedelta) -> None: 

96 """Moves the clock forwards, or backwards if delta is negative.""" 

97 self._refuse_mid_transaction() 

98 self.frozen_at += delta 

99 

100 def _db_settings(self) -> tuple[str, str]: 

101 return "", self.frozen_at.isoformat() 

102 

103 

104def _check_aware(when: datetime) -> None: 

105 if when.tzinfo is None: 

106 raise ValueError("timewarp needs an aware datetime, this one has no timezone") 

107 

108 

109# mock goes ahead of pg_catalog so that an unqualified now() finds mock.now() first 

110MOCK_SEARCH_PATH = "public, mock, pg_catalog" 

111 

112 

113def create_mock_clock(conn: Connection) -> None: 

114 """ 

115 Installs mock.now(), the postgres half of this fixture: the frozen instant if there is one, and 

116 otherwise postgres' own clock displaced by the offset. Both settings are read per transaction, 

117 and set by install_timewarp below. 

118 

119 This lives outside create_schema_from_models because it has to exist before *either* way of 

120 building the schema runs, and survive drop_database() in between, so that migrations and 

121 models bake the same function into their column defaults. 

122 """ 

123 conn.execute( 

124 text(""" 

125 CREATE SCHEMA mock; 

126 

127 CREATE FUNCTION mock.now() RETURNS timestamptz 

128 LANGUAGE sql STABLE AS $$ 

129 SELECT coalesce( 

130 nullif(current_setting('mock.frozen_at', true), '')::timestamptz, 

131 pg_catalog.now() + coalesce( 

132 nullif(current_setting('mock.offset', true), '')::interval, 

133 interval '0' 

134 ) 

135 ); 

136 $$; 

137 """) 

138 ) 

139 conn.commit() 

140 

141 

142def mock_clock_installed(conn: Connection) -> bool: 

143 """ 

144 Whether this database has the mock clock, ie whether shifting it does anything at all. In one 

145 built without it, an unqualified now() is still pg_catalog.now() and reports the real time. 

146 """ 

147 installed: bool = conn.exec_driver_sql( 

148 "SELECT to_regprocedure('now()') IS NOT DISTINCT FROM to_regprocedure('mock.now()')" 

149 ).scalar_one() 

150 return installed 

151 

152 

153def _as_interval(delta: timedelta) -> str: 

154 # these three are exact and sum to delta, unlike total_seconds() which is a float 

155 return f"{delta.days} days {delta.seconds} seconds {delta.microseconds} microseconds" 

156 

157 

158def install_timewarp[WarpT: Clock](warp: WarpT) -> Generator[WarpT]: 

159 # the engine is only connected to once something runs a query, so this is fine without the db fixture 

160 engine = _get_base_engine() 

161 

162 # postgres reads the clock through mock.now(), which returns mock.frozen_at if there is one and 

163 # otherwise adds mock.offset to its own clock. Both are set per transaction so they can't leak 

164 # into another test via a pooled connection. 

165 def sync_db_clock(conn: Connection) -> None: 

166 conn.exec_driver_sql( 

167 "SELECT set_config('mock.offset', %s, true), set_config('mock.frozen_at', %s, true)", 

168 warp._db_settings(), 

169 ) 

170 if not mock_clock_installed(conn): 

171 raise RuntimeError( 

172 "timewarp can't move the postgres clock in this database: now() still resolves to " 

173 "pg_catalog.now(), so the database would quietly report the real time. mock.now() " 

174 "is installed when the test database is built, so request the `db` fixture in any " 

175 "test that uses timewarp and touches the database." 

176 ) 

177 warp._open_transactions.add(conn) 

178 

179 def transaction_ended(conn: Connection) -> None: 

180 warp._open_transactions.discard(conn) 

181 

182 event.listen(engine, "begin", sync_db_clock) 

183 event.listen(engine, "commit", transaction_ended) 

184 event.listen(engine, "rollback", transaction_ended) 

185 try: 

186 with patch.object(utils, "_mockable_now", warp.now): 

187 yield warp 

188 finally: 

189 event.remove(engine, "begin", sync_db_clock) 

190 event.remove(engine, "commit", transaction_ended) 

191 event.remove(engine, "rollback", transaction_ended)