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

@@ -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