diff --git a/superset/utils/screenshot_utils.py b/superset/utils/screenshot_utils.py index 82a06fd66c0c..290e22dd3647 100644 --- a/superset/utils/screenshot_utils.py +++ b/superset/utils/screenshot_utils.py @@ -281,6 +281,7 @@ def take_tiled_screenshot( # noqa: C901 load_wait: int = 60, animation_wait: int = 0, log_context: str | None = None, + screenshot_started_at: float | None = None, ) -> bytes | None: """ Take a tiled screenshot of a large dashboard by scrolling and capturing sections. @@ -294,6 +295,11 @@ def take_tiled_screenshot( # noqa: C901 log_context: Optional identifier (e.g. report execution id, or a cache key for thumbnails) appended to log lines so a slow/timed-out capture can be traced back to the run that produced it. + screenshot_started_at: Optional time.monotonic() timestamp taken at + the start of the overall screenshot operation (before browser + navigation), so time already spent on goto/headstart/element + waits counts against the budget -- the same clock the non-tiled + readiness wait uses. Falls back to "now" when not provided. Returns: Combined screenshot bytes or None if failed @@ -323,7 +329,15 @@ def take_tiled_screenshot( # noqa: C901 wait_budget_seconds = resolve_screenshot_task_budget_seconds(log_context) if wait_budget_seconds is None: wait_budget_seconds = float(TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS) - start_time = time.monotonic() + # Anchor the budget clock at the start of the overall screenshot + # operation when the caller provides it: the budget is derived from the + # Celery task's own limit, and navigation/headstart/element waits before + # tiling spend that same task time. A locally-reset clock would grant + # the tile loop a fresh full budget on top of whatever was already + # spent, which on hard-limit-only tasks can still crest the task limit. + start_time = ( + screenshot_started_at if screenshot_started_at is not None else time.monotonic() + ) try: # Get the target element element = page.locator(f".{element_name}") diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py index 4b8c18372754..05e856be4138 100644 --- a/superset/utils/webdriver.py +++ b/superset/utils/webdriver.py @@ -592,6 +592,7 @@ def get_screenshot( # pylint: disable=too-many-locals, too-many-statements # n load_wait=self._screenshot_load_wait, animation_wait=selenium_animation_wait, log_context=log_context, + screenshot_started_at=screenshot_started_at, ) if not img: # _get_screenshot() has no wait/readiness logic at diff --git a/tests/unit_tests/utils/test_screenshot_utils.py b/tests/unit_tests/utils/test_screenshot_utils.py index f5569e8a7419..9c25379a7df2 100644 --- a/tests/unit_tests/utils/test_screenshot_utils.py +++ b/tests/unit_tests/utils/test_screenshot_utils.py @@ -915,6 +915,64 @@ def test_derived_task_budget_caps_tile_wait(self, mock_page): assert first_timeout == 96 * 1000 assert first_timeout < 200 * 1000 + def test_screenshot_started_at_counts_pre_capture_time_against_budget( + self, mock_page, monkeypatch + ): + """Time spent before tiling (navigation, headstart, element waits) + must count against the budget when the caller provides the overall + screenshot start time -- the budget is derived from the Celery + task's own limit, and that same task time was already being spent + before the tile loop started. Same clock the non-tiled readiness + wait uses via _wait_for_charts_ready.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + # The overall screenshot started 900s ago (goto/headstart/element + # waits consumed it); only 100s of the 1000s budget remains. + clock.now = 900.0 + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + load_wait=500, + screenshot_started_at=0.0, + ) + + first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"] + assert first_timeout == 100 * 1000 + + def test_omitted_screenshot_started_at_anchors_clock_locally( + self, mock_page, monkeypatch + ): + """Without the caller-provided anchor the clock starts at entry + (backward-compatible default): the same 900s of prior wall-clock + time does not reduce the budget.""" + monkeypatch.setattr( + "superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS", # noqa: E501 + 1000, + ) + clock = self._FakeClock() + clock.now = 900.0 + + with patch("superset.utils.screenshot_utils.current_task", None): + with patch("superset.utils.screenshot_utils.time.monotonic", new=clock): + with patch("superset.utils.screenshot_utils.combine_screenshot_tiles"): + take_tiled_screenshot( + mock_page, + "dashboard", + tile_height=2000, + load_wait=500, + ) + + first_timeout = mock_page.wait_for_function.call_args_list[0][1]["timeout"] + assert first_timeout == 500 * 1000 + def test_fast_dashboard_matches_default_behavior(self, mock_page): """Well under budget, waits are not capped and behavior is unchanged.""" with patch("superset.utils.screenshot_utils.current_task", None): diff --git a/tests/unit_tests/utils/webdriver_test.py b/tests/unit_tests/utils/webdriver_test.py index 8b0ab64c3642..73da0ec609fd 100644 --- a/tests/unit_tests/utils/webdriver_test.py +++ b/tests/unit_tests/utils/webdriver_test.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from unittest.mock import MagicMock, patch, PropertyMock +from unittest.mock import ANY, MagicMock, patch, PropertyMock import pytest @@ -1681,6 +1681,7 @@ def test_tiled_path_passes_animation_wait_per_tile_no_global_wait( load_wait=30, animation_wait=2, log_context=None, + screenshot_started_at=ANY, ) # The only wait_for_timeout call should be the 0ms headstart; no global # animation wait should be issued (handled per-tile by take_tiled_screenshot)