1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Tests for screenshot cache bug fixes:
1. Cache only saved when image generation succeeds
2. Recompute stale COMPUTING tasks and UPDATED without image
"""
from datetime import datetime, timedelta
from unittest.mock import MagicMock, patch
import pytest
from pytest_mock import MockerFixture
from superset.utils.screenshots import (
BaseScreenshot,
ScreenshotCachePayload,
StatusValues,
)
BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot"
class MockCache:
"""A class to manage screenshot cache for testing."""
def __init__(self):
self._cache = {}
def set(self, key, value):
"""Set the cache with a new value."""
self._cache[key] = value
def get(self, key):
"""Get the cached value."""
return self._cache.get(key)
def clear(self):
"""Clear all cached values."""
self._cache.clear()
@pytest.fixture
def mock_user():
"""Fixture to create a mock user."""
user = MagicMock()
user.id = 1
return user
@pytest.fixture
def screenshot_obj():
"""Fixture to create a BaseScreenshot object."""
url = "http://example.com"
digest = "sample_digest"
return BaseScreenshot(url, digest)
class TestCacheOnlyOnSuccess:
"""Test that cache is only saved when image generation succeeds."""
def _setup_mocks(self, mocker: MockerFixture, screenshot_obj):
"""Helper method to set up common mocks."""
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
get_screenshot = mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"image_data"
)
# Mock resize_image to avoid PIL errors with fake image data
mocker.patch(
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
)
BaseScreenshot.cache = MockCache()
return get_screenshot
def test_cache_error_status_when_screenshot_fails(
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""Test that error status is cached when screenshot generation fails."""
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
get_screenshot = mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
side_effect=Exception("Screenshot failed"),
)
BaseScreenshot.cache = MockCache()
# Execute compute_and_cache
screenshot_obj.compute_and_cache(user=mock_user, force=True)
# Verify get_screenshot was called
get_screenshot.assert_called_once()
# Cache should be set with ERROR status (to prevent immediate retries)
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Error"
assert cached_value.get("image") is None
def test_cache_error_status_when_resize_fails(
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""Test that error status is cached when image resize fails."""
self._setup_mocks(mocker, screenshot_obj)
mocker.patch(
BASE_SCREENSHOT_PATH + ".resize_image",
side_effect=Exception("Resize failed"),
)
# Use different window and thumb sizes to trigger resize
screenshot_obj.compute_and_cache(
user=mock_user, force=True, window_size=(800, 600), thumb_size=(400, 300)
)
# Cache should be set with ERROR status (to prevent immediate retries)
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Error"
assert cached_value.get("image") is None
def test_cache_saved_only_when_image_generated(
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""Test that cache is only saved when image is successfully generated."""
self._setup_mocks(mocker, screenshot_obj)
# Execute compute_and_cache
screenshot_obj.compute_and_cache(user=mock_user, force=True)
# Cache should be set with UPDATED status and image
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Updated"
assert cached_value["image"] is not None
def test_no_intermediate_cache_during_computing(
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""Test that cache is not saved during COMPUTING state."""
mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
BaseScreenshot.cache = MockCache()
# Mock get_screenshot to check cache state during execution
def check_cache_during_screenshot(*args, **kwargs):
# At this point, we're in COMPUTING state
# Cache should NOT be set yet
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
# Cache should be empty during screenshot generation
assert cached_value is None, (
"Cache should not be saved during COMPUTING state"
)
return b"image_data"
mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
side_effect=check_cache_during_screenshot,
)
# Mock resize to avoid PIL errors with fake image data
mocker.patch(
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image_data"
)
# Execute compute_and_cache
screenshot_obj.compute_and_cache(user=mock_user, force=True)
# After completion, cache should be set with UPDATED status
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Updated"
class TestShouldTriggerTask:
"""Test the should_trigger_task method improvements."""
@patch("superset.utils.screenshots.app")
def test_trigger_on_stale_computing_status(self, mock_app):
"""Test that stale COMPUTING status triggers recomputation."""
# Set TTL to 300 seconds
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Create payload with COMPUTING status from 400 seconds ago (stale)
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=old_timestamp
)
# Should trigger task because COMPUTING is stale
assert payload.should_trigger_task(force=False) is True
@patch("superset.utils.screenshots.app")
def test_no_trigger_on_fresh_computing_status(self, mock_app):
"""Test that fresh COMPUTING status does not trigger recomputation."""
# Set TTL to 300 seconds
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Create payload with COMPUTING status from 100 seconds ago (fresh)
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=fresh_timestamp
)
# Should NOT trigger task because COMPUTING is still fresh
assert payload.should_trigger_task(force=False) is False
def test_trigger_on_updated_without_image(self):
"""Test that UPDATED status without image triggers recomputation."""
# Create payload with UPDATED status but no image
# This simulates the bug where cache was saved without an image
payload = ScreenshotCachePayload(image=None, status=StatusValues.UPDATED)
# Should trigger task because UPDATED but has no image
assert payload.should_trigger_task(force=False) is True
def test_no_trigger_on_updated_with_image(self):
"""Test that UPDATED status with image does not trigger recomputation."""
# Create payload with UPDATED status and valid image
payload = ScreenshotCachePayload(image=b"valid_image_data")
# Should NOT trigger task because UPDATED with valid image
assert payload.should_trigger_task(force=False) is False
def test_trigger_on_pending_status(self):
"""Test that PENDING status triggers task."""
payload = ScreenshotCachePayload(status=StatusValues.PENDING)
assert payload.should_trigger_task(force=False) is True
@patch("superset.utils.screenshots.app")
def test_trigger_on_expired_error(self, mock_app):
"""Test that expired ERROR status triggers task."""
# Set TTL to 300 seconds
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Create payload with ERROR status from 400 seconds ago (expired)
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.ERROR, timestamp=old_timestamp
)
assert payload.should_trigger_task(force=False) is True
@patch("superset.utils.screenshots.app")
def test_no_trigger_on_fresh_error(self, mock_app):
"""Test that fresh ERROR status does not trigger task."""
# Set TTL to 300 seconds
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Create payload with ERROR status from 100 seconds ago (fresh)
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.ERROR, timestamp=fresh_timestamp
)
assert payload.should_trigger_task(force=False) is False
def test_force_always_triggers(self):
"""Test that force=True always triggers task regardless of status."""
# Test with UPDATED + image (normally wouldn't trigger)
payload_updated = ScreenshotCachePayload(image=b"image_data")
assert payload_updated.should_trigger_task(force=True) is True
# Test with fresh COMPUTING (normally wouldn't trigger)
payload_computing = ScreenshotCachePayload(status=StatusValues.COMPUTING)
assert payload_computing.should_trigger_task(force=True) is True
class TestIsComputingStale:
"""Test the is_computing_stale method."""
@patch("superset.utils.screenshots.app")
def test_computing_is_stale(self, mock_app):
"""Test that old COMPUTING status is detected as stale."""
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Timestamp from 400 seconds ago
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=old_timestamp
)
assert payload.is_computing_stale() is True
@patch("superset.utils.screenshots.app")
def test_computing_is_not_stale(self, mock_app):
"""Test that fresh COMPUTING status is not stale."""
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Timestamp from 100 seconds ago
fresh_timestamp = (datetime.now() - timedelta(seconds=100)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=fresh_timestamp
)
assert payload.is_computing_stale() is False
@patch("superset.utils.screenshots.app")
def test_computing_exactly_at_ttl(self, mock_app):
"""Test boundary condition at exactly TTL."""
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Timestamp from exactly 300 seconds ago
exact_timestamp = (datetime.now() - timedelta(seconds=300)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=exact_timestamp
)
# At exactly TTL, should be stale (>= TTL)
assert payload.is_computing_stale() is True
@patch("superset.utils.screenshots.app")
def test_computing_just_past_ttl(self, mock_app):
"""Test boundary condition just past TTL."""
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
# Timestamp from 301 seconds ago (just past TTL)
past_ttl_timestamp = (datetime.now() - timedelta(seconds=301)).isoformat()
payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=past_ttl_timestamp
)
# Just past TTL should be stale
assert payload.is_computing_stale() is True
class TestIntegrationCacheBugFix:
"""Integration tests combining both fixes."""
def test_failed_screenshot_does_not_pollute_cache(
self, mocker: MockerFixture, screenshot_obj, mock_user
):
"""
Integration test: Failed screenshot should cache error status
to prevent immediate retries, not leave corrupted cache with image=None.
"""
mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot",
side_effect=Exception("Network error"),
)
BaseScreenshot.cache = MockCache()
# First attempt fails
screenshot_obj.compute_and_cache(user=mock_user, force=True)
# Verify cache contains ERROR status (prevents immediate retry)
cache_key = screenshot_obj.get_cache_key()
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Error"
assert cached_value.get("image") is None
# Cache entry should not trigger task immediately (error is fresh)
cached_payload = screenshot_obj.get_from_cache_key(cache_key)
assert cached_payload is not None
assert cached_payload.should_trigger_task(force=False) is False
@patch("superset.utils.screenshots.app")
def test_stale_computing_triggers_retry(
self, mock_app, mocker: MockerFixture, screenshot_obj, mock_user
):
"""
Integration test: Stale COMPUTING status should trigger retry
to recover from stuck tasks.
"""
mock_app.config = {"THUMBNAIL_ERROR_CACHE_TTL": 300}
BaseScreenshot.cache = MockCache()
# Create stale COMPUTING entry and seed it in the cache
old_timestamp = (datetime.now() - timedelta(seconds=400)).isoformat()
stale_payload = ScreenshotCachePayload(
status=StatusValues.COMPUTING, timestamp=old_timestamp
)
cache_key = screenshot_obj.get_cache_key()
BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
mocker.patch(
BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=b"recovered_image"
)
# Mock resize to avoid PIL errors
mocker.patch(
BASE_SCREENSHOT_PATH + ".resize_image", return_value=b"resized_image"
)
# Should trigger task because COMPUTING is stale
assert stale_payload.should_trigger_task() is True
# Retry should succeed and update cache
screenshot_obj.compute_and_cache(user=mock_user, force=False)
cached_value = BaseScreenshot.cache.get(cache_key)
assert cached_value is not None
assert cached_value["status"] == "Updated"
assert cached_value["image"] is not None