Skip to content
Closed
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
16 changes: 15 additions & 1 deletion superset/utils/screenshot_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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}")
Expand Down
1 change: 1 addition & 0 deletions superset/utils/webdriver.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
58 changes: 58 additions & 0 deletions tests/unit_tests/utils/test_screenshot_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
3 changes: 2 additions & 1 deletion tests/unit_tests/utils/webdriver_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Expand Down
Loading