Coverage for ivatar/test_access_stats.py: 98%

1import base64 

2from django.test import TransactionTestCase, Client, override_settings 

3from django.contrib.auth.models import User 

4from ivatar.ivataraccount.models import ( 

5 ConfirmedEmail, 

6 Photo, 

7 ConfirmedEmailAccessStat, 

8) 

9from ivatar.access_stats import stats_manager 

10from django.core.cache import cache 

11import time 

12from unittest.mock import patch 

13 

14 

15@override_settings( 

16 CACHES={ 

17 "default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}, 

18 "filesystem": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}, 

19 }, 

20 CACHE_RESPONSE=False, 

21 STATS_BATCH_SIZE=100, 

22 STATS_FLUSH_TIMEOUT=5, # Short timeout for testing 

23 ASYNC_ACCESS_COUNT=True, 

24) 

25class BatchStatsTest(TransactionTestCase): 

26 def setUp(self): 

27 self.username = "testuser_" + str(time.time()) 

28 self.password = "password" 

29 self.user = User.objects.create_user( 

30 username=self.username, password=self.password 

31 ) 

32 self.client = Client() 

33 

34 # Create a confirmed email and photo 

35 self.email_address = f"test_{time.time()}@example.com" 

36 self.email = ConfirmedEmail.objects.create( 

37 user=self.user, email=self.email_address, ip_address="127.0.0.1" 

38 ) 

39 

40 # Create a dummy photo (minimal PNG, Base64-encoded for pycco compatibility) 

41 png_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAACXBIWXMAAAsTAAALEwEAmpwYAAAACklEQVQIHWNgYAAAAAMAAWveibQAAAAASUVORK5CYII=" 

42 png_data = base64.b64decode(png_base64) 

43 self.photo = Photo.objects.create( 

44 user=self.user, data=png_data, format="png", ip_address="127.0.0.1" 

45 ) 

46 self.email.photo = self.photo 

47 self.email.save() 

48 

49 # Clear keys for this test 

50 model_name = self.email.__class__.__name__ 

51 pk = self.email.pk 

52 self.email_key = f"stats:access_count:{model_name}:{pk}" 

53 self.start_time_key = f"stats:start_time:{model_name}:{pk}" 

54 

55 cache.delete(self.email_key) 

56 cache.delete(self.start_time_key) 

57 

58 # Ensure stats_manager has clean state 

59 stats_manager._local_dirty_keys.clear() 

60 

61 # Mock threading.Thread to run synchronously 

62 self.thread_patcher = patch("ivatar.access_stats.threading.Thread") 

63 self.mock_thread_cls = self.thread_patcher.start() 

64 

65 # Set up side effect to run target immediately 

66 def side_effect(target=None, args=(), **kwargs): 

67 class SyncThread: 

68 def start(self): 

69 if target: 

70 target(*args) 

71 

72 return SyncThread() 

73 

74 self.mock_thread_cls.side_effect = side_effect 

75 

76 def tearDown(self): 

77 self.thread_patcher.stop() 

78 

79 def _get_email_access_count(self): 

80 """Read access_count from the stat table for self.email.""" 

81 try: 

82 return ConfirmedEmailAccessStat.objects.get(email=self.email).access_count 

83 except ConfirmedEmailAccessStat.DoesNotExist: 

84 return 0 

85 

86 @override_settings(ASYNC_ACCESS_COUNT=False) 

87 def test_sync_update_when_disabled(self): 

88 """ 

89 Verify that updates are synchronous when feature flag is disabled. 

90 """ 

91 stats_manager.update_access_count(self.email) 

92 

93 # Check stat table updated immediately 

94 self.assertEqual(self._get_email_access_count(), 1) 

95 

96 # Check cache should be empty (or at least not used for this logic) 

97 self.assertIsNone(cache.get(self.email_key)) 

98 

99 def test_normal_increment_cached(self): 

100 """ 

101 Verify that a single request increments cache but does not update DB immediately. 

102 """ 

103 stats_manager.update_access_count(self.email) 

104 

105 # Check cache 

106 self.assertEqual(cache.get(self.email_key), 1) 

107 

108 # Check stat table (should be 0) 

109 self.assertEqual(self._get_email_access_count(), 0) 

110 

111 def test_batch_flush(self): 

112 """ 

113 Verify that hitting the batch limit triggers a DB update. 

114 """ 

115 # Increment 99 times (limit is 100) 

116 for _ in range(99): 

117 stats_manager.update_access_count(self.email) 

118 

119 self.assertEqual(cache.get(self.email_key), 99) 

120 self.assertEqual(self._get_email_access_count(), 0) 

121 

122 # 100th increment -> Triggers flush (synchronously due to mock) 

123 stats_manager.update_access_count(self.email) 

124 

125 # No wait needed 

126 

127 # Check stat table 

128 self.assertEqual(self._get_email_access_count(), 100) 

129 

130 # Check cache reset 

131 self.assertEqual(cache.get(self.email_key), 0) 

132 

133 def test_time_flush(self): 

134 """ 

135 Verify that the time-based flush works. 

136 Timeout set to 5 seconds in override_settings. 

137 """ 

138 # 1. First increment sets start time 

139 stats_manager.update_access_count(self.email) 

140 

141 # Verify start time is set 

142 self.assertIsNotNone(cache.get(self.start_time_key)) 

143 

144 # Verify no DB update yet 

145 self.assertEqual(self._get_email_access_count(), 0) 

146 

147 # 2. Wait for timeout > 5s 

148 time.sleep(6) 

149 

150 # 3. Next increment should trigger flush due to timeout (synchronously) 

151 stats_manager.update_access_count(self.email) 

152 

153 # No wait needed 

154 

155 # Check stat table: should have 2 updates (1 from before wait + 1 from trigger) 

156 self.assertEqual(self._get_email_access_count(), 2) 

157 

158 # Check cache reset 

159 self.assertEqual(cache.get(self.email_key), 0) 

160 # Start time should be cleared 

161 self.assertIsNone(cache.get(self.start_time_key)) 

162 

163 def test_atexit_handler(self): 

164 """ 

165 Verify that the cleanup handler flushes remaining counts. 

166 """ 

167 # Add some counts 

168 stats_manager.update_access_count(self.email) 

169 stats_manager.update_access_count(self.email) 

170 stats_manager.update_access_count(self.email) 

171 

172 self.assertEqual(cache.get(self.email_key), 3) 

173 self.assertEqual(self._get_email_access_count(), 0) 

174 

175 # Manually call the handler 

176 stats_manager.flush_all_dirty_keys() 

177 

178 # Check stat table 

179 self.assertEqual(self._get_email_access_count(), 3) 

180 

181 # Check cache decremented 

182 self.assertEqual(cache.get(self.email_key), 0)