This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch enforce-max-screenshot-width in repository https://gitbox.apache.org/repos/asf/superset.git
commit 90e9328c1377cde9b3598213f436b6dea5772761 Author: Beto Dealmeida <[email protected]> AuthorDate: Fri Jan 31 10:30:06 2025 -0500 Add tests --- tests/unit_tests/commands/report/execute_test.py | 115 +++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/tests/unit_tests/commands/report/execute_test.py b/tests/unit_tests/commands/report/execute_test.py index 3d49bb0457..7c04f2f8ab 100644 --- a/tests/unit_tests/commands/report/execute_test.py +++ b/tests/unit_tests/commands/report/execute_test.py @@ -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 @@ -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 + + [email protected]( + "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]}" + )
