fix: enforce `ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH` (#32053)

This commit is contained in:
Beto Dealmeida 2025-01-31 14:56:56 -05:00 committed by GitHub
parent 7db0589340
commit 1064ad5d58
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 130 additions and 11 deletions

View File

@ -300,13 +300,16 @@ class BaseReportState:
) )
user = security_manager.find_user(username) user = security_manager.find_user(username)
max_width = app.config["ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH"]
if self._report_schedule.chart: if self._report_schedule.chart:
url = self._get_url() url = self._get_url()
window_width, window_height = app.config["WEBDRIVER_WINDOW"]["slice"] window_width, window_height = app.config["WEBDRIVER_WINDOW"]["slice"]
window_size = ( width = min(max_width, self._report_schedule.custom_width or window_width)
self._report_schedule.custom_width or window_width, height = self._report_schedule.custom_height or window_height
self._report_schedule.custom_height or window_height, window_size = (width, height)
)
screenshots: list[Union[ChartScreenshot, DashboardScreenshot]] = [ screenshots: list[Union[ChartScreenshot, DashboardScreenshot]] = [
ChartScreenshot( ChartScreenshot(
url, url,
@ -317,11 +320,12 @@ class BaseReportState:
] ]
else: else:
urls = self.get_dashboard_urls() urls = self.get_dashboard_urls()
window_width, window_height = app.config["WEBDRIVER_WINDOW"]["dashboard"] window_width, window_height = app.config["WEBDRIVER_WINDOW"]["dashboard"]
window_size = ( width = min(max_width, self._report_schedule.custom_width or window_width)
self._report_schedule.custom_width or window_width, height = self._report_schedule.custom_height or window_height
self._report_schedule.custom_height or window_height, window_size = (width, height)
)
screenshots = [ screenshots = [
DashboardScreenshot( DashboardScreenshot(
url, url,
@ -578,9 +582,9 @@ class BaseReportState:
SupersetError( SupersetError(
message=ex.message, message=ex.message,
error_type=SupersetErrorType.REPORT_NOTIFICATION_ERROR, error_type=SupersetErrorType.REPORT_NOTIFICATION_ERROR,
level=ErrorLevel.ERROR level=(
if ex.status >= 500 ErrorLevel.ERROR if ex.status >= 500 else ErrorLevel.WARNING
else ErrorLevel.WARNING, ),
) )
) )
if notification_errors: if notification_errors:

View File

@ -16,19 +16,24 @@
# under the License. # under the License.
import json import json
from datetime import datetime
from unittest.mock import patch from unittest.mock import patch
from uuid import UUID
import pytest import pytest
from pytest_mock import MockerFixture from pytest_mock import MockerFixture
from superset.app import SupersetApp
from superset.commands.report.execute import BaseReportState from superset.commands.report.execute import BaseReportState
from superset.dashboards.permalink.types import DashboardPermalinkState from superset.dashboards.permalink.types import DashboardPermalinkState
from superset.reports.models import ( from superset.reports.models import (
ReportRecipientType, ReportRecipientType,
ReportSchedule, ReportSchedule,
ReportScheduleType,
ReportSourceFormat, ReportSourceFormat,
) )
from superset.utils.core import HeaderDataType from superset.utils.core import HeaderDataType
from superset.utils.screenshots import ChartScreenshot
from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.conftest import with_feature_flags
@ -365,3 +370,113 @@ def test_get_tab_url(
) )
result: str = class_instance._get_tab_url(dashboard_state) result: str = class_instance._get_tab_url(dashboard_state)
assert result == "http://0.0.0.0:8080/superset/dashboard/p/uri/" assert result == "http://0.0.0.0:8080/superset/dashboard/p/uri/"
def create_report_schedule(
mocker: MockerFixture,
custom_width: int | None = None,
custom_height: int | None = None,
) -> ReportSchedule:
"""Helper function to create a ReportSchedule instance with specified dimensions."""
schedule = ReportSchedule()
schedule.type = ReportScheduleType.REPORT
schedule.name = "Test Report"
schedule.description = "Test Description"
schedule.chart = mocker.MagicMock()
schedule.chart.id = 1
schedule.dashboard = None
schedule.database = None
schedule.custom_width = custom_width
schedule.custom_height = custom_height
return schedule
@pytest.mark.parametrize(
"test_id,custom_width,max_width,window_width,expected_width",
[
# Test when custom width exceeds max width
("exceeds_max", 2000, 1600, 800, 1600),
# Test when custom width is less than max width
("under_max", 1200, 1600, 800, 1200),
# Test when custom width is None (should use window width)
("no_custom", None, 1600, 800, 800),
# Test when custom width equals max width
("equals_max", 1600, 1600, 800, 1600),
],
)
def test_screenshot_width_calculation(
app: SupersetApp,
mocker: MockerFixture,
test_id: str,
custom_width: int | None,
max_width: int,
window_width: int,
expected_width: int,
) -> None:
"""
Test that screenshot width is correctly calculated.
The width should be:
- Limited by max_width when custom_width exceeds it
- Equal to custom_width when it's less than max_width
- Equal to window_width when custom_width is None
"""
from superset.commands.report.execute import BaseReportState
# Mock configuration
app.config.update(
{
"ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH": max_width,
"WEBDRIVER_WINDOW": {
"slice": (window_width, 600),
"dashboard": (window_width, 600),
},
"ALERT_REPORTS_EXECUTORS": {},
}
)
# Create report schedule with specified custom width
report_schedule = create_report_schedule(mocker, custom_width=custom_width)
# Initialize BaseReportState
report_state = BaseReportState(
report_schedule=report_schedule,
scheduled_dttm=datetime.now(),
execution_id=UUID("084e7ee6-5557-4ecd-9632-b7f39c9ec524"),
)
# Mock security manager and screenshot
with (
patch(
"superset.commands.report.execute.security_manager"
) as mock_security_manager,
patch(
"superset.utils.screenshots.ChartScreenshot.get_screenshot"
) as mock_get_screenshot,
):
# Mock user
mock_user = mocker.MagicMock()
mock_security_manager.find_user.return_value = mock_user
mock_get_screenshot.return_value = b"screenshot bytes"
# Mock get_executor to avoid database lookups
with patch(
"superset.commands.report.execute.get_executor"
) as mock_get_executor:
mock_get_executor.return_value = ("executor", "username")
# Capture the ChartScreenshot instantiation
with patch(
"superset.commands.report.execute.ChartScreenshot",
wraps=ChartScreenshot,
) as mock_chart_screenshot:
# Call the method that triggers screenshot creation
report_state._get_screenshots()
# Verify ChartScreenshot was created with correct window_size
mock_chart_screenshot.assert_called_once()
_, kwargs = mock_chart_screenshot.call_args
assert kwargs["window_size"][0] == expected_width, (
f"Test {test_id}: Expected width {expected_width}, "
f"but got {kwargs['window_size'][0]}"
)