fix(thumbnail cache): Enabling force parameter on screenshot/thumbnail cache (#31757)

Co-authored-by: Kamil Gabryjelski <kamil.gabryjelski@gmail.com>
This commit is contained in:
Jack
2025-01-31 12:22:31 -06:00
committed by GitHub
parent c590e90c87
commit 7db0589340
13 changed files with 609 additions and 280 deletions

View File

@@ -319,9 +319,5 @@ def test_compute_thumbnails(thumbnail_mock, app_context, fs):
["-d", "-i", dashboard.id],
)
thumbnail_mock.assert_called_with(
None,
dashboard.id,
force=False,
)
thumbnail_mock.assert_called_with(None, dashboard.id, force=False)
assert response.exit_code == 0

View File

@@ -37,6 +37,7 @@ from superset.reports.models import ReportSchedule, ReportScheduleType
from superset.models.slice import Slice
from superset.tags.models import Tag, TaggedObject, TagType, ObjectType
from superset.utils.core import backend, override_user
from superset.utils.screenshots import ScreenshotCachePayload
from superset.utils import json
from tests.integration_tests.base_api_tests import ApiOwnersTestCaseMixin
@@ -3069,13 +3070,15 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
@pytest.mark.usefixtures("create_dashboard_with_tag")
@patch("superset.dashboards.api.cache_dashboard_screenshot")
@patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
def test_screenshot_success_png(self, mock_get_cache, mock_cache_task):
def test_screenshot_success_png(self, mock_get_from_cache_key, mock_cache_task):
"""
Validate screenshot returns png
"""
self.login(ADMIN_USERNAME)
mock_cache_task.return_value = None
mock_get_cache.return_value = BytesIO(b"fake image data")
mock_get_from_cache_key.return_value = ScreenshotCachePayload(
b"fake image data"
)
dashboard = (
db.session.query(Dashboard)
@@ -3083,7 +3086,7 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
.first()
)
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
assert cache_resp.status_code == 200
cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
response = self._get_screenshot(dashboard.id, cache_key, "png")
@@ -3091,20 +3094,29 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
assert response.mimetype == "image/png"
assert response.data == b"fake image data"
mock_get_from_cache_key.return_value = ScreenshotCachePayload()
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
@with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
@pytest.mark.usefixtures("create_dashboard_with_tag")
@patch("superset.dashboards.api.cache_dashboard_screenshot")
@patch("superset.dashboards.api.build_pdf_from_screenshots")
@patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
def test_screenshot_success_pdf(
self, mock_get_from_cache, mock_build_pdf, mock_cache_task
self,
mock_get_from_cache_key,
mock_build_pdf,
mock_cache_task,
):
"""
Validate screenshot can return pdf.
"""
self.login(ADMIN_USERNAME)
mock_cache_task.return_value = None
mock_get_from_cache.return_value = BytesIO(b"fake image data")
mock_get_from_cache_key.return_value = ScreenshotCachePayload(
b"fake image data"
)
mock_build_pdf.return_value = b"fake pdf data"
dashboard = (
@@ -3113,7 +3125,7 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
.first()
)
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
assert cache_resp.status_code == 200
cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
response = self._get_screenshot(dashboard.id, cache_key, "pdf")
@@ -3121,6 +3133,10 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
assert response.mimetype == "application/pdf"
assert response.data == b"fake pdf data"
mock_get_from_cache_key.return_value = ScreenshotCachePayload()
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
@with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True)
@pytest.mark.usefixtures("create_dashboard_with_tag")
@patch("superset.dashboards.api.cache_dashboard_screenshot")
@@ -3153,10 +3169,12 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
@pytest.mark.usefixtures("create_dashboard_with_tag")
@patch("superset.dashboards.api.cache_dashboard_screenshot")
@patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key")
def test_screenshot_invalid_download_format(self, mock_get_cache, mock_cache_task):
def test_screenshot_invalid_download_format(
self, mock_get_from_cache_key, mock_cache_task
):
self.login(ADMIN_USERNAME)
mock_cache_task.return_value = None
mock_get_cache.return_value = BytesIO(b"fake png data")
mock_get_from_cache_key.return_value = ScreenshotCachePayload(b"fake png data")
dashboard = (
db.session.query(Dashboard)
@@ -3165,9 +3183,13 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas
)
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
assert cache_resp.status_code == 200
cache_key = json.loads(cache_resp.data.decode("utf-8"))["cache_key"]
mock_get_from_cache_key.return_value = ScreenshotCachePayload()
cache_resp = self._cache_screenshot(dashboard.id)
assert cache_resp.status_code == 202
response = self._get_screenshot(dashboard.id, cache_key, "invalid")
assert response.status_code == 404

View File

@@ -18,7 +18,6 @@
# from superset.models.dashboard import Dashboard
import urllib.request
from io import BytesIO
from unittest import skipUnless
from unittest.mock import ANY, call, MagicMock, patch
@@ -32,7 +31,11 @@ from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
from superset.tasks.types import ExecutorType, FixedExecutor
from superset.utils import json
from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
from superset.utils.screenshots import (
ChartScreenshot,
DashboardScreenshot,
ScreenshotCachePayload,
)
from superset.utils.urls import get_url_path
from superset.utils.webdriver import WebDriverSelenium
from tests.integration_tests.base_tests import SupersetTestCase
@@ -287,14 +290,14 @@ class TestThumbnails(SupersetTestCase):
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
@skipUnless((is_feature_enabled("THUMBNAILS")), "Thumbnails feature")
def test_get_async_dashboard_not_allowed(self):
def test_get_async_dashboard_created(self):
"""
Thumbnails: Simple get async dashboard not allowed
"""
self.login(ADMIN_USERNAME)
_, thumbnail_url = self._get_id_and_thumbnail_url(DASHBOARD_URL)
rv = self.client.get(thumbnail_url)
assert rv.status_code == 404
assert rv.status_code == 202
@pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
@with_feature_flags(THUMBNAILS=True)
@@ -370,7 +373,9 @@ class TestThumbnails(SupersetTestCase):
Thumbnails: Simple get chart with wrong digest
"""
with patch.object(
ChartScreenshot, "get_from_cache", return_value=BytesIO(self.mock_image)
ChartScreenshot,
"get_from_cache",
return_value=ScreenshotCachePayload(self.mock_image),
):
self.login(ADMIN_USERNAME)
id_, thumbnail_url = self._get_id_and_thumbnail_url(CHART_URL)
@@ -385,7 +390,9 @@ class TestThumbnails(SupersetTestCase):
Thumbnails: Simple get cached dashboard screenshot
"""
with patch.object(
DashboardScreenshot, "get_from_cache", return_value=BytesIO(self.mock_image)
DashboardScreenshot,
"get_from_cache_key",
return_value=ScreenshotCachePayload(self.mock_image),
):
self.login(ADMIN_USERNAME)
_, thumbnail_url = self._get_id_and_thumbnail_url(DASHBOARD_URL)
@@ -400,7 +407,9 @@ class TestThumbnails(SupersetTestCase):
Thumbnails: Simple get cached chart screenshot
"""
with patch.object(
ChartScreenshot, "get_from_cache", return_value=BytesIO(self.mock_image)
ChartScreenshot,
"get_from_cache_key",
return_value=ScreenshotCachePayload(self.mock_image),
):
self.login(ADMIN_USERNAME)
id_, thumbnail_url = self._get_id_and_thumbnail_url(CHART_URL)
@@ -415,7 +424,9 @@ class TestThumbnails(SupersetTestCase):
Thumbnails: Simple get dashboard with wrong digest
"""
with patch.object(
DashboardScreenshot, "get_from_cache", return_value=BytesIO(self.mock_image)
DashboardScreenshot,
"get_from_cache",
return_value=ScreenshotCachePayload(self.mock_image),
):
self.login(ADMIN_USERNAME)
id_, thumbnail_url = self._get_id_and_thumbnail_url(DASHBOARD_URL)