You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2025/08/21 18:47:29 UTC
(superset) 01/01: Add tests
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch get-image-always-return
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 85681f42ecd0c6924e0e770a30129a670c5a1e26
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Thu Aug 21 14:30:47 2025 -0400
Add tests
---
superset/charts/api.py | 13 ++-
superset/dashboards/api.py | 12 +-
superset/exceptions.py | 4 +
superset/utils/screenshots.py | 5 +-
tests/unit_tests/utils/screenshot_test.py | 47 ++++++++
.../utils/test_screenshot_exception_handling.py | 126 +++++++++++++++++++++
6 files changed, 201 insertions(+), 6 deletions(-)
diff --git a/superset/charts/api.py b/superset/charts/api.py
index a33c0bc943..66c0781ffb 100644
--- a/superset/charts/api.py
+++ b/superset/charts/api.py
@@ -80,6 +80,7 @@ from superset.commands.importers.exceptions import (
from superset.commands.importers.v1.utils import get_contents_from_bundle
from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
from superset.daos.chart import ChartDAO
+from superset.exceptions import ScreenshotImageNotAvailableException
from superset.extensions import event_logger
from superset.models.slice import Slice
from superset.tasks.thumbnails import cache_chart_thumbnail
@@ -704,8 +705,12 @@ class ChartRestApi(BaseSupersetModelRestApi):
if cache_payload := ChartScreenshot.get_from_cache_key(digest):
if cache_payload.status == StatusValues.UPDATED:
+ try:
+ image = cache_payload.get_image()
+ except ScreenshotImageNotAvailableException:
+ return self.response_404()
return Response(
- FileWrapper(cache_payload.get_image()),
+ FileWrapper(image),
mimetype="image/png",
direct_passthrough=True,
)
@@ -791,8 +796,12 @@ class ChartRestApi(BaseSupersetModelRestApi):
task_status=cache_payload.get_status(),
)
self.incr_stats("from_cache", self.thumbnail.__name__)
+ try:
+ image = cache_payload.get_image()
+ except ScreenshotImageNotAvailableException:
+ return self.response_404()
return Response(
- FileWrapper(cache_payload.get_image()),
+ FileWrapper(image),
mimetype="image/png",
direct_passthrough=True,
)
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 43bbfc1ec5..a739e44c27 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -101,6 +101,7 @@ from superset.dashboards.schemas import (
TabsPayloadSchema,
thumbnail_query_schema,
)
+from superset.exceptions import ScreenshotImageNotAvailableException
from superset.extensions import event_logger
from superset.models.dashboard import Dashboard
from superset.models.embedded_dashboard import EmbeddedDashboard
@@ -1199,7 +1200,10 @@ class DashboardRestApi(BaseSupersetModelRestApi):
# fetch the dashboard screenshot using the current user and cache if set
if cache_payload := DashboardScreenshot.get_from_cache_key(digest):
- image = cache_payload.get_image()
+ try:
+ image = cache_payload.get_image()
+ except ScreenshotImageNotAvailableException:
+ return self.response_404()
if download_format == "pdf":
pdf_img = image.getvalue()
# Convert the screenshot to PDF
@@ -1322,8 +1326,12 @@ class DashboardRestApi(BaseSupersetModelRestApi):
)
self.incr_stats("from_cache", self.thumbnail.__name__)
+ try:
+ image = cache_payload.get_image()
+ except ScreenshotImageNotAvailableException:
+ return self.response_404()
return Response(
- FileWrapper(cache_payload.get_image()),
+ FileWrapper(image),
mimetype="image/png",
direct_passthrough=True,
)
diff --git a/superset/exceptions.py b/superset/exceptions.py
index 6dddfb8302..1b3216da0c 100644
--- a/superset/exceptions.py
+++ b/superset/exceptions.py
@@ -486,3 +486,7 @@ class SupersetResultsBackendNotConfigureException(SupersetErrorException):
level=ErrorLevel.ERROR,
)
super().__init__(error)
+
+
+class ScreenshotImageNotAvailableException(SupersetException):
+ status = 404
diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py
index 250cd89a6d..e11253014b 100644
--- a/superset/utils/screenshots.py
+++ b/superset/utils/screenshots.py
@@ -23,9 +23,10 @@ from enum import Enum
from io import BytesIO
from typing import cast, TYPE_CHECKING, TypedDict
-from flask import abort, current_app as app
+from flask import current_app as app
from superset import feature_flag_manager, thumbnail_cache
+from superset.exceptions import ScreenshotImageNotAvailableException
from superset.extensions import event_logger
from superset.utils.hashing import md5_sha_from_dict
from superset.utils.urls import modify_url_query
@@ -123,7 +124,7 @@ class ScreenshotCachePayload:
def get_image(self) -> BytesIO:
if self._image is None:
- abort(404)
+ raise ScreenshotImageNotAvailableException()
return BytesIO(cast(bytes, self._image))
def get_timestamp(self) -> str:
diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py
index fa928fc73f..d747a90055 100644
--- a/tests/unit_tests/utils/screenshot_test.py
+++ b/tests/unit_tests/utils/screenshot_test.py
@@ -192,3 +192,50 @@ class TestComputeAndCache:
force=False, window_size=(1, 1), thumb_size=thumb_size
)
resize_image.assert_called_once()
+
+
+class TestScreenshotCachePayloadGetImage:
+ """Test the get_image method behavior including exception handling"""
+
+ def test_get_image_returns_bytesio_when_image_exists(self):
+ """Test that get_image returns BytesIO object when image data exists"""
+ image_data = b"test image data"
+ payload = ScreenshotCachePayload(image=image_data)
+
+ result = payload.get_image()
+
+ assert result is not None
+ assert result.read() == image_data
+
+ def test_get_image_raises_exception_when_no_image(self):
+ """Test get_image raises ScreenshotImageNotAvailableException when no image"""
+ from superset.exceptions import ScreenshotImageNotAvailableException
+
+ payload = ScreenshotCachePayload() # No image data
+
+ with pytest.raises(ScreenshotImageNotAvailableException):
+ payload.get_image()
+
+ def test_get_image_raises_exception_when_image_is_none(self):
+ """Test that get_image raises exception when image is explicitly set to None"""
+ from superset.exceptions import ScreenshotImageNotAvailableException
+
+ payload = ScreenshotCachePayload(image=None)
+
+ with pytest.raises(ScreenshotImageNotAvailableException):
+ payload.get_image()
+
+ def test_get_image_multiple_reads(self):
+ """Test that get_image returns fresh BytesIO each time"""
+ image_data = b"test image data"
+ payload = ScreenshotCachePayload(image=image_data)
+
+ result1 = payload.get_image()
+ result2 = payload.get_image()
+
+ # Both should be valid BytesIO objects
+ assert result1.read() == image_data
+ assert result2.read() == image_data
+
+ # Should be different BytesIO instances
+ assert result1 is not result2
diff --git a/tests/unit_tests/utils/test_screenshot_exception_handling.py b/tests/unit_tests/utils/test_screenshot_exception_handling.py
new file mode 100644
index 0000000000..4d1f6bee7b
--- /dev/null
+++ b/tests/unit_tests/utils/test_screenshot_exception_handling.py
@@ -0,0 +1,126 @@
+# 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 exception handling in API endpoints.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from superset.exceptions import ScreenshotImageNotAvailableException
+from superset.utils.screenshots import ScreenshotCachePayload
+
+
+class TestScreenshotAPIExceptionHandling:
+ """Test that API endpoints properly handle ScreenshotImageNotAvailableException"""
+
+ def test_dashboard_screenshot_api_handles_exception(self):
+ """Test dashboard screenshot API returns 404 when get_image raises exception"""
+ from superset.dashboards.api import DashboardRestApi
+
+ # Create mock API instance
+ api = DashboardRestApi()
+ api.response_404 = MagicMock(return_value="404 Response")
+
+ # Create mock cache payload that will raise exception
+ mock_cache_payload = MagicMock()
+ mock_cache_payload.get_image.side_effect = (
+ ScreenshotImageNotAvailableException()
+ )
+
+ # Test the exception handling logic
+ try:
+ mock_cache_payload.get_image()
+ pytest.fail("Should have raised exception")
+ except ScreenshotImageNotAvailableException:
+ response = api.response_404()
+ assert response == "404 Response"
+
+ def test_chart_screenshot_api_handles_exception(self):
+ """Test that chart screenshot API returns 404 when get_image raises exception"""
+ from superset.charts.api import ChartRestApi
+
+ # Create mock API instance
+ api = ChartRestApi()
+ api.response_404 = MagicMock(return_value="404 Response")
+
+ # Create mock cache payload that will raise exception
+ mock_cache_payload = MagicMock()
+ mock_cache_payload.get_image.side_effect = (
+ ScreenshotImageNotAvailableException()
+ )
+
+ # Test the exception handling logic
+ try:
+ mock_cache_payload.get_image()
+ pytest.fail("Should have raised exception")
+ except ScreenshotImageNotAvailableException:
+ response = api.response_404()
+ assert response == "404 Response"
+
+ def test_screenshot_cache_payload_exception_has_correct_status(self):
+ """Test that the ScreenshotImageNotAvailableException has status 404"""
+ exception = ScreenshotImageNotAvailableException()
+ assert exception.status == 404
+
+ def test_api_method_simulation_with_exception(self):
+ """Simulate the API method behavior with exception handling"""
+
+ def simulate_dashboard_screenshot_method(cache_payload):
+ """Simulate the logic in dashboard screenshot methods"""
+ try:
+ image = cache_payload.get_image()
+ return {"status": "success", "image": image}
+ except ScreenshotImageNotAvailableException:
+ return {"status": "404", "message": "Not Found"}
+
+ # Test with payload that has image
+ payload_with_image = ScreenshotCachePayload(image=b"test data")
+ result = simulate_dashboard_screenshot_method(payload_with_image)
+ assert result["status"] == "success"
+ assert result["image"] is not None
+
+ # Test with payload that has no image (should raise exception)
+ payload_no_image = ScreenshotCachePayload()
+ result = simulate_dashboard_screenshot_method(payload_no_image)
+ assert result["status"] == "404"
+ assert result["message"] == "Not Found"
+
+ def test_api_method_simulation_with_file_wrapper(self):
+ """Simulate the FileWrapper usage in API methods"""
+ from werkzeug.wsgi import FileWrapper
+
+ def simulate_api_file_response(cache_payload):
+ """Simulate the FileWrapper logic in API methods"""
+ try:
+ image = cache_payload.get_image()
+ return FileWrapper(image)
+ except ScreenshotImageNotAvailableException:
+ return None
+
+ # Test with valid image
+ payload_with_image = ScreenshotCachePayload(image=b"test data")
+ result = simulate_api_file_response(payload_with_image)
+ assert result is not None
+ assert isinstance(result, FileWrapper)
+
+ # Test without image
+ payload_no_image = ScreenshotCachePayload()
+ result = simulate_api_file_response(payload_no_image)
+ assert result is None