bito-code-review[bot] commented on code in PR #42118:
URL: https://github.com/apache/superset/pull/42118#discussion_r3635509152
##########
tests/unit_tests/utils/test_screenshot_utils.py:
##########
@@ -568,3 +576,279 @@ def test_animation_wait_default_is_zero(self):
sig = inspect.signature(take_tiled_screenshot)
assert sig.parameters["animation_wait"].default == 0
+
+
+class TestTileWaitBudget:
+ @pytest.fixture
+ def mock_page(self):
+ """Create a mock Playwright page object for a 3-tile (5000px)
dashboard."""
+ page = MagicMock()
+ element = MagicMock()
+ page.locator.return_value = element
+ page.evaluate.return_value = {
+ "height": 5000,
+ "top": 100,
+ "left": 50,
+ "width": 800,
+ }
+ page.screenshot.return_value = b"fake_screenshot_data"
+ return page
+
+ class _FakeClock:
+ """Stateful monotonic() stand-in the test advances explicitly.
+
+ Robust to how many times the code under test samples the clock per
+ tile (budget check, per-tile wait timing, animation budget) -- only
+ explicit advances move time forward.
+ """
+
+ def __init__(self) -> None:
+ self.now = 0.0
+
+ def __call__(self) -> float:
+ return self.now
+
+ def test_per_tile_wait_shrinks_as_budget_depletes(self, mock_page,
monkeypatch):
+ """Each tile's readiness-wait timeout is capped at the remaining
budget."""
+ monkeypatch.setattr(
+
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS",
# noqa: E501
+ 1000,
+ )
+ clock = self._FakeClock()
+ # Simulate slow tiles: the readiness wait itself consumes wall time,
+ # so each subsequent tile sees less remaining budget.
+ wait_durations = iter([950, 40, 5])
+
+ def slow_wait(*args, **kwargs):
+ clock.now += next(wait_durations)
+
+ mock_page.wait_for_function.side_effect = slow_wait
+
+ with patch("superset.utils.screenshot_utils.time.monotonic",
new=clock):
+ with
patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
+ result = take_tiled_screenshot(
+ mock_page, "dashboard", tile_height=2000, load_wait=100
+ )
+
+ assert result is not None
+ timeouts = [
+ call[1]["timeout"] for call in
mock_page.wait_for_function.call_args_list
+ ]
+ # remaining budget at each tile's wait: 1000, 50, 10 seconds
+ # -> capped timeouts shrink
+ assert timeouts == [100 * 1000, 50 * 1000, 10 * 1000]
+ assert timeouts == sorted(timeouts, reverse=True)
+
+ def test_budget_exhausted_raises_and_stops_capturing(self, mock_page,
monkeypatch):
+ """Exhausting the budget aborts cleanly instead of capturing
unchecked."""
+ monkeypatch.setattr(
+
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS",
# noqa: E501
+ 1000,
+ )
+ clock = self._FakeClock()
+ # Tile 0's readiness wait consumes the whole budget; tile 1's budget
+ # check then sees remaining <= 0 and raises before capturing.
+ mock_page.wait_for_function.side_effect = lambda *args, **kwargs:
setattr(
+ clock, "now", 1000.0
+ )
+
+ with patch("superset.utils.screenshot_utils.time.monotonic",
new=clock):
+ with patch(
+ "superset.utils.screenshot_utils.combine_screenshot_tiles"
+ ) as mock_combine:
+ with patch("superset.utils.screenshot_utils.logger") as
mock_logger:
+ with pytest.raises(TiledScreenshotBudgetExceededError):
+ take_tiled_screenshot(
+ mock_page, "dashboard", tile_height=2000,
load_wait=100
+ )
+
+ # Only the first tile was captured before the budget ran out.
+ assert mock_page.screenshot.call_count == 1
+ # Tiles were never combined -- the function raised before that point.
+ mock_combine.assert_not_called()
+
+ # Budget exhaustion is a customer chart-loading issue, not a Superset
+ # system fault, so it must log at WARNING (not ERROR) -- consistent
+ # with the #38130/#38441 precedent for screenshot timeout logging.
+ assert mock_logger.error.call_count == 0
+ mock_logger.warning.assert_called_once()
+ warning_args = mock_logger.warning.call_args[0]
+ assert "budget exhausted" in warning_args[0]
+ # tile index, tiles total, tiles captured, tiles total,
+ # elapsed seconds, budget seconds, log-context suffix
+ assert warning_args[1] == 2
+ assert warning_args[2] == 3
+ assert warning_args[3] == 1
+ assert warning_args[4] == 3
+ assert warning_args[5] == 1000
+ assert warning_args[6] == 1000
+ assert warning_args[7] == ""
+
+ def test_budget_exhausted_warning_includes_log_context(
+ self, mock_page, monkeypatch
+ ):
+ """log_context (e.g. report execution id) is appended to the
warning."""
+ monkeypatch.setattr(
+
"superset.utils.screenshot_utils.TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS",
# noqa: E501
+ 1000,
+ )
+ clock = self._FakeClock()
+ # Tile 0's readiness wait consumes the whole budget; tile 1's budget
+ # check then sees remaining <= 0 and raises.
+ mock_page.wait_for_function.side_effect = lambda *args, **kwargs:
setattr(
+ clock, "now", 1000.0
+ )
+ with patch("superset.utils.screenshot_utils.time.monotonic",
new=clock):
+ with
patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
+ with patch("superset.utils.screenshot_utils.logger") as
mock_logger:
+ with pytest.raises(TiledScreenshotBudgetExceededError):
+ take_tiled_screenshot(
+ mock_page,
+ "dashboard",
+ tile_height=2000,
+ load_wait=100,
+ log_context="execution_id=abc-123",
+ )
+
+ warning_args = mock_logger.warning.call_args[0]
+ assert warning_args[-1] == " [execution_id=abc-123]"
+
+ def test_per_tile_timing_debug_line_logged(self, mock_page):
+ """Each tile logs a DEBUG timing breakdown (spinner wait, animation
wait)."""
+ with patch("superset.utils.screenshot_utils.logger") as mock_logger:
+ with
patch("superset.utils.screenshot_utils.combine_screenshot_tiles"):
+ take_tiled_screenshot(
+ mock_page,
+ "dashboard",
+ tile_height=2000,
+ log_context="cache_key=xyz",
+ )
+
+ timing_calls = [
+ call for call in mock_logger.debug.call_args_list if "timing" in
call[0][0]
+ ]
+ assert len(timing_calls) == 3
+ for i, call in enumerate(timing_calls):
+ args = call[0]
+ assert args[1] == i + 1 # tile index
+ assert args[2] == 3 # total tiles
+ assert args[-1] == " [cache_key=xyz]"
+
+ 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.combine_screenshot_tiles"):
+ result = take_tiled_screenshot(
+ mock_page,
+ "dashboard",
+ tile_height=2000,
+ load_wait=30,
+ animation_wait=5,
+ )
+
+ assert result is not None
+ assert mock_page.screenshot.call_count == 3
+
+ for call in mock_page.wait_for_function.call_args_list:
+ assert call[1]["timeout"] == 30 * 1000
+
+ animation_calls = [
+ call
+ for call in mock_page.wait_for_timeout.call_args_list
+ if call[0][0] == 5 * 1000
+ ]
+ assert len(animation_calls) == 3
+
+
+class TestResolveWaitBudgetSeconds:
+ """The budget is derived from the running Celery task's own time limit
+ when available, and falls back to the fixed constant otherwise."""
+
+ def _mock_task(self, soft=None, hard=None):
+ task = MagicMock()
+ task.request.timelimit = (soft, hard)
+ return task
+
+ def test_derives_budget_from_soft_time_limit(self):
+ """soft_time_limit is preferred over the hard time_limit when both are
set."""
+ task = self._mock_task(soft=90, hard=120)
+ with patch("superset.utils.screenshot_utils.current_task", task):
+ budget = _resolve_wait_budget_seconds()
+
+ # margin = min(300, 90 * 0.2) = 18; budget = 90 - 18 = 72
+ assert budget == 72
+
+ def test_small_task_limit_yields_positive_scaled_margin_budget(self):
+ """A 120s thumbnail-task limit (superset-shell#4389) still gets a
+ usable, positive budget via the scaled-down margin, not the fixed
+ 300s margin that would otherwise wipe it out."""
+ task = self._mock_task(soft=None, hard=120)
+ with patch("superset.utils.screenshot_utils.current_task", task):
+ budget = _resolve_wait_budget_seconds()
+
+ # margin = min(300, 120 * 0.2) = 24; budget = 120 - 24 = 96
+ assert budget == 96
+ assert budget > 0
+ assert budget < 120
+
+ def test_no_task_context_falls_back_to_constant(self):
+ """Outside of a Celery task, the fixed fallback budget is used."""
+ with patch("superset.utils.screenshot_utils.current_task", None):
+ budget = _resolve_wait_budget_seconds()
+
+ assert budget == TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS
+
+ def test_task_with_no_timelimit_falls_back_to_constant(self):
+ """A task with no soft or hard limit set falls back to the constant."""
+ task = self._mock_task(soft=None, hard=None)
+ with patch("superset.utils.screenshot_utils.current_task", task):
+ budget = _resolve_wait_budget_seconds()
+
+ assert budget == TILED_SCREENSHOT_TOTAL_WAIT_BUDGET_SECONDS
+
+ def test_derivation_exception_falls_back_to_constant(self):
+ """Any failure while inspecting the task context must never break a
+ screenshot -- fall back to the constant and log at DEBUG."""
+
+ class _BrokenTask:
+ """Simulates a task-like object whose .request raises."""
+
+ @property
+ def request(self):
+ raise RuntimeError("boom")
+
+ with patch("superset.utils.screenshot_utils.current_task",
_BrokenTask()):
Review Comment:
<div>
<div id="suggestion">
<div id="issue"><b>Exception uses string literal directly</b></div>
<div id="fix">
Exception raised with string literal directly instead of assigning to a
variable first.
</div>
<details>
<summary>
<b>Code suggestion</b>
</summary>
<blockquote>Check the AI-generated fix before applying</blockquote>
<div id="code">
````suggestion
def request(self):
error_msg = "boom"
raise RuntimeError(error_msg)
with patch("superset.utils.screenshot_utils.current_task",
_BrokenTask()):
````
</div>
</details>
</div>
<small><i>Code Review Run #1c7538</i></small>
</div>
---
Should Bito avoid suggestions like this for future reviews? (<a
href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>)
- [ ] Yes, avoid them
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
