Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: enforce ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH #32053

Merged
merged 2 commits into from
Jan 31, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 15 additions & 11 deletions superset/commands/report/execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -300,13 +300,16 @@ def _get_screenshots(self) -> list[bytes]:
)
user = security_manager.find_user(username)

max_width = app.config["ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH"]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@betodealmeida Maybe we should deprecate the ALERT_REPORTS_MAX_CUSTOM_SCREENSHOT_WIDTH feature flag to remove it in a major version. It seems that this feature should always be enabled in case alerts and reports are too.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, it's not a feature, just a configuration value.


if self._report_schedule.chart:
url = self._get_url()

window_width, window_height = app.config["WEBDRIVER_WINDOW"]["slice"]
window_size = (
self._report_schedule.custom_width or window_width,
self._report_schedule.custom_height or window_height,
)
width = min(max_width, self._report_schedule.custom_width or window_width)
height = self._report_schedule.custom_height or window_height
window_size = (width, height)

screenshots: list[Union[ChartScreenshot, DashboardScreenshot]] = [
ChartScreenshot(
url,
Expand All @@ -317,11 +320,12 @@ def _get_screenshots(self) -> list[bytes]:
]
else:
urls = self.get_dashboard_urls()

window_width, window_height = app.config["WEBDRIVER_WINDOW"]["dashboard"]
window_size = (
self._report_schedule.custom_width or window_width,
self._report_schedule.custom_height or window_height,
)
width = min(max_width, self._report_schedule.custom_width or window_width)
height = self._report_schedule.custom_height or window_height
window_size = (width, height)

screenshots = [
DashboardScreenshot(
url,
Expand Down Expand Up @@ -578,9 +582,9 @@ def _send(
SupersetError(
message=ex.message,
error_type=SupersetErrorType.REPORT_NOTIFICATION_ERROR,
level=ErrorLevel.ERROR
if ex.status >= 500
else ErrorLevel.WARNING,
level=(
ErrorLevel.ERROR if ex.status >= 500 else ErrorLevel.WARNING
),
)
)
if notification_errors:
Expand Down
115 changes: 115 additions & 0 deletions tests/unit_tests/commands/report/execute_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,19 +16,24 @@
# under the License.

import json
from datetime import datetime
from unittest.mock import patch
from uuid import UUID

import pytest
from pytest_mock import MockerFixture

from superset.app import SupersetApp
from superset.commands.report.execute import BaseReportState
from superset.dashboards.permalink.types import DashboardPermalinkState
from superset.reports.models import (
ReportRecipientType,
ReportSchedule,
ReportScheduleType,
ReportSourceFormat,
)
from superset.utils.core import HeaderDataType
from superset.utils.screenshots import ChartScreenshot
from tests.integration_tests.conftest import with_feature_flags


Expand Down Expand Up @@ -365,3 +370,113 @@ def test_get_tab_url(
)
result: str = class_instance._get_tab_url(dashboard_state)
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(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice!

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]}"
)
Loading