refactor(screenshots): return ScreenshotCaptureResult from capture paths

Introduces a frozen ScreenshotCaptureResult dataclass (image, readiness)
and threads it through the driver and screenshot capture layers in
place of a raw bytes | None return type: WebDriverProxy.get_screenshot,
WebDriverPlaywright.get_screenshot, BaseScreenshot.get_screenshot, and
their five call sites. Behavior-preserving -- readiness is always None
here; a later change narrows that field's type once a readiness report
type exists. Call sites now read capture_result.image instead of
treating the return value's truthiness directly.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

superset-frontend/src/components/Chart/Chart.test.tsx

@@ -18,11 +18,34 @@
  */
 import { render, screen, fireEvent } from 'spec/helpers/testing-library';
 import '@testing-library/jest-dom';
+import type { JsonObject } from '@superset-ui/core';
 import { PLACEHOLDER_DATASOURCE } from 'src/dashboard/constants';
 import { ResourceStatus } from 'src/hooks/apiResources/apiResources';
+import type { ChartStatus } from 'src/explore/types';
+import {
+  CHART_STATUS_ATTR,
+  CHART_STATUS_ATTR_VALUES,
+} from 'src/utils/screenshotContract';
 import Chart from './Chart';
 import type { Actions } from './Chart';
 
+// The success/rendered/pending branches render down through ChartRenderer
+// into SuperChart, which needs a registered viz plugin. Mock it the same way
+// ChartRenderer.test.tsx does, so these tests can exercise the beacon
+// attribute on Chart's own root node without depending on the chart plugin
+// registry.
+jest.mock('@superset-ui/core', () => ({
+  ...jest.requireActual('@superset-ui/core'),
+  SuperChart: (props: JsonObject) => (
+    <div data-test="mock-super-chart">{JSON.stringify(props.formData)}</div>
+  ),
+}));
+
+jest.mock(
+  'src/components/Chart/ChartContextMenu/ChartContextMenu',
+  () => () => <div data-test="mock-chart-context-menu" />,
+);
+
 const mockActions: Actions = {
   logEvent: jest.fn() as unknown as Actions['logEvent'],
   chartRenderingFailed: jest.fn() as unknown as Actions['chartRenderingFailed'],
@@ -41,55 +64,108 @@ const baseProps = {
   setControlValue: jest.fn(),
 };
 
-test('shows backend error instead of loading spinner when datasource is still a placeholder', () => {
-  render(
-    <Chart
-      {...baseProps}
-      chartStatus="failed"
-      chartAlert="Your default credentials were not found."
-      datasource={PLACEHOLDER_DATASOURCE}
-      datasetsStatus={ResourceStatus.Loading}
-      queriesResponse={[
+// Finds the root DOM node carrying the data-chart-status beacon attribute,
+// regardless of which Chart.tsx render branch produced it.
+const getStatusNode = (container: HTMLElement) =>
+  container.querySelector(`[${CHART_STATUS_ATTR}]`);
+
+// dashboardId is intentionally varied across these cases (rather than only
+// tested once via statusCases below) because it changes which sub-tree
+// actually renders for "failed" (spinner size, ChartSource passed into
+// ChartErrorMessage) and for "stopped" (EmptyState sizing) — the published
+// beacon value must stay identical across both render paths regardless.
+const failedBackendErrorProps = {
+  chartStatus: 'failed' as const,
+  chartAlert: 'Your default credentials were not found.',
+  datasource: PLACEHOLDER_DATASOURCE,
+  datasetsStatus: ResourceStatus.Loading,
+  queriesResponse: [
+    {
+      errors: [
         {
-          errors: [
-            {
-              error_type: 'GENERIC_BACKEND_ERROR',
-              message: 'Your default credentials were not found.',
-              extra: {
-                issue_codes: [{ code: 1011, message: 'Issue 1011' }],
-              },
-              level: 'error',
-            },
-          ],
+          error_type: 'GENERIC_BACKEND_ERROR',
+          message: 'Your default credentials were not found.',
+          extra: {
+            issue_codes: [{ code: 1011, message: 'Issue 1011' }],
+          },
+          level: 'error',
         },
-      ]}
-    />,
+      ],
+    },
+  ],
+};
+
+test('shows backend error instead of loading spinner when datasource is still a placeholder (standalone)', () => {
+  const { container } = render(
+    <Chart {...baseProps} {...failedBackendErrorProps} />,
   );
 
   expect(
     screen.getByText(/Your default credentials were not found/),
   ).toBeInTheDocument();
+  // A real backend error is a terminal failure, so the beacon must say
+  // "failed", not "loading" — even though the datasource is still a
+  // placeholder, the errors array takes precedence.
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.FAILED,
+  );
 });
 
-test('shows loading spinner for client-side errors without errors array when datasource is still a placeholder', () => {
-  render(
-    <Chart
-      {...baseProps}
-      chartStatus="failed"
-      chartAlert="Some client-side error"
-      datasource={PLACEHOLDER_DATASOURCE}
-      datasetsStatus={ResourceStatus.Loading}
-      queriesResponse={[{}]}
-    />,
+test('shows backend error instead of loading spinner when datasource is still a placeholder (dashboard)', () => {
+  const { container } = render(
+    <Chart {...baseProps} {...failedBackendErrorProps} dashboardId={7} />,
+  );
+
+  expect(
+    screen.getByText(/Your default credentials were not found/),
+  ).toBeInTheDocument();
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.FAILED,
+  );
+});
+
+const failedDatasourceLoadingProps = {
+  chartStatus: 'failed' as const,
+  chartAlert: 'Some client-side error',
+  datasource: PLACEHOLDER_DATASOURCE,
+  datasetsStatus: ResourceStatus.Loading,
+  queriesResponse: [{}],
+};
+
+test('shows loading spinner for client-side errors without errors array when datasource is still a placeholder (standalone)', () => {
+  const { container } = render(
+    <Chart {...baseProps} {...failedDatasourceLoadingProps} />,
   );
 
   expect(screen.getByRole('status')).toBeInTheDocument();
   expect(screen.queryByText(/Some client-side error/)).not.toBeInTheDocument();
+  // Datasource-still-loading sub-case of the failed branch: publishing the
+  // terminal "failed" value here would lie to beacon consumers about a
+  // spinner state that is still in progress.
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.LOADING,
+  );
 });
 
-test('shows the stop message and a re-run affordance when the query was stopped', () => {
+test('shows loading spinner for client-side errors without errors array when datasource is still a placeholder (dashboard)', () => {
+  const { container } = render(
+    <Chart {...baseProps} {...failedDatasourceLoadingProps} dashboardId={7} />,
+  );
+
+  expect(screen.getByRole('status')).toBeInTheDocument();
+  expect(screen.queryByText(/Some client-side error/)).not.toBeInTheDocument();
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.LOADING,
+  );
+});
+
+test('shows the stop message and a re-run affordance when the query was stopped (standalone)', () => {
   const onQuery = jest.fn();
-  render(
+  const { container } = render(
     <Chart
       {...baseProps}
       chartStatus="stopped"
@@ -104,4 +180,156 @@ test('shows the stop message and a re-run affordance when the query was stopped'
   expect(rerun.tagName).toBe('BUTTON');
   fireEvent.click(rerun);
   expect(onQuery).toHaveBeenCalledTimes(1);
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.STOPPED,
+  );
+});
+
+test('shows the stop message and a re-run affordance when the query was stopped (dashboard)', () => {
+  const onQuery = jest.fn();
+  const { container } = render(
+    <Chart
+      {...baseProps}
+      dashboardId={7}
+      chartStatus="stopped"
+      chartAlert="Updating chart was stopped"
+      onQuery={onQuery}
+    />,
+  );
+
+  expect(screen.getByText('Updating chart was stopped')).toBeInTheDocument();
+
+  const rerun = screen.getByText('click here');
+  expect(rerun.tagName).toBe('BUTTON');
+  fireEvent.click(rerun);
+  expect(onQuery).toHaveBeenCalledTimes(1);
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.STOPPED,
+  );
+});
+
+// Every reducer-set chartStatus value, plus the pre-fetch `null` that Redux
+// starts with, mapped to the beacon value Chart.tsx must publish on its
+// shared Styles root — for both the standalone (Explore) and dashboard
+// render paths, since dashboardId changes which sub-tree renders (spinner
+// size, muted styling) but must never change the published status.
+const statusCases: Array<{
+  chartStatus: ChartStatus | '' | null | undefined;
+  expected: string;
+  label: string;
+}> = [
+  {
+    chartStatus: null,
+    expected: CHART_STATUS_ATTR_VALUES.PENDING,
+    label: 'null',
+  },
+  {
+    chartStatus: '',
+    expected: CHART_STATUS_ATTR_VALUES.PENDING,
+    label: '"" (datasource-retry reset)',
+  },
+  {
+    chartStatus: 'loading',
+    expected: CHART_STATUS_ATTR_VALUES.LOADING,
+    label: '"loading"',
+  },
+  {
+    chartStatus: 'success',
+    expected: CHART_STATUS_ATTR_VALUES.LOADING,
+    label: '"success" (pre-paint, non-terminal)',
+  },
+  {
+    chartStatus: 'rendered',
+    expected: CHART_STATUS_ATTR_VALUES.RENDERED,
+    label: '"rendered"',
+  },
+];
+
+statusCases.forEach(({ chartStatus, expected, label }) => {
+  test(`publishes data-chart-status="${expected}" for chartStatus ${label} (standalone)`, () => {
+    const { container } = render(
+      <Chart
+        {...baseProps}
+        chartStatus={chartStatus as ChartStatus}
+        queriesResponse={[]}
+      />,
+    );
+    expect(screen.getByTestId('chart-container')).toHaveAttribute(
+      CHART_STATUS_ATTR,
+      expected,
+    );
+    expect(getStatusNode(container)).toBe(
+      screen.getByTestId('chart-container'),
+    );
+  });
+
+  test(`publishes data-chart-status="${expected}" for chartStatus ${label} (dashboard)`, () => {
+    const { container } = render(
+      <Chart
+        {...baseProps}
+        dashboardId={7}
+        chartStatus={chartStatus as ChartStatus}
+        queriesResponse={[]}
+      />,
+    );
+    expect(screen.getByTestId('chart-container')).toHaveAttribute(
+      CHART_STATUS_ATTR,
+      expected,
+    );
+    expect(getStatusNode(container)).toBe(
+      screen.getByTestId('chart-container'),
+    );
+  });
+});
+
+test('falls back to data-chart-status="pending" for an unrecognized/future chartStatus value', () => {
+  const { container } = render(
+    <Chart
+      {...baseProps}
+      chartStatus={'not-a-real-status' as unknown as ChartStatus}
+      queriesResponse={[]}
+    />,
+  );
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.PENDING,
+  );
+});
+
+test('publishes data-chart-status="stopped" for the missing-controls EmptyState (Explore preview, not a chartStatus)', () => {
+  const { container } = render(
+    <Chart {...baseProps} errorMessage="Add required control values" />,
+  );
+  expect(
+    screen.getByText('Add required control values to preview chart'),
+  ).toBeInTheDocument();
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.STOPPED,
+  );
+});
+
+test('publishes data-chart-status="stopped" for the "ready to go" EmptyState (Explore preview, not a chartStatus)', () => {
+  const { container } = render(
+    <Chart {...baseProps} chartIsStale queriesResponse={[]} />,
+  );
+  expect(screen.getByText('Your chart is ready to go!')).toBeInTheDocument();
+  expect(getStatusNode(container)).toHaveAttribute(
+    CHART_STATUS_ATTR,
+    CHART_STATUS_ATTR_VALUES.STOPPED,
+  );
+});
+
+test('screenshotContract exports no CSS class name strings', () => {
+  // Guard against scope creep: this contract is attribute-only, and must
+  // never grow a styling hook such as a CSS class name.
+  expect(CHART_STATUS_ATTR).toBe('data-chart-status');
+  Object.values(CHART_STATUS_ATTR_VALUES).forEach(value => {
+    expect(typeof value).toBe('string');
+    expect(value).not.toMatch(/[.#]/); // no CSS selector punctuation
+    expect(value.toLowerCase()).not.toContain('chart-container');
+    expect(value.toLowerCase()).not.toContain('class');
+  });
 });

superset-frontend/src/components/Chart/Chart.tsx

@@ -46,6 +46,12 @@ import { Dispatch } from 'redux';
 import ChartRenderer from './ChartRenderer';
 import { ChartErrorMessage } from './ChartErrorMessage';
 import { getChartRequiredFieldsMissingMessage } from '../../utils/getChartRequiredFieldsMissingMessage';
+import {
+  CHART_STATUS_ATTR,
+  CHART_STATUS_ATTR_VALUES,
+  getChartStatusAttrValue,
+  type ChartStatusAttrValue,
+} from '../../utils/screenshotContract';
 
 export type ChartErrorType = Partial<ClientErrorObject>;
 export interface ChartProps {
@@ -126,6 +132,13 @@ const NONEXISTENT_DATASET = t(
   'The dataset associated with this chart no longer exists',
 );
 
+// JSX prop names must be static identifiers, so the CHART_STATUS_ATTR
+// constant (the single source of truth for the attribute name) is bound via
+// a computed-key spread rather than a literal `data-chart-status=` prop.
+const chartStatusAttrProps = (value: ChartStatusAttrValue) => ({
+  [CHART_STATUS_ATTR]: value,
+});
+
 const Styles = styled.div<{ height: number; width?: number }>`
   min-height: ${p => p.height}px;
   position: relative;
@@ -430,8 +443,30 @@ function Chart({
   const showSpinner = isLoading && !suppressLoadingSpinner;
 
   if (chartStatus === 'failed') {
+    // If none of the query responses carry a backend error, renderErrorMessage
+    // falls back to a "datasource still loading" spinner instead of an error
+    // message (see the condition mirrored below). In that sub-case, the
+    // terminal "failed" value would misrepresent an in-progress state to
+    // consumers of this attribute, so publish "loading" instead.
+    const hasQueryErrors = ensureIsArray(queriesResponse).some(
+      item => (item as ChartErrorType)?.errors?.[0],
+    );
+    const isDatasourceStillLoading =
+      !hasQueryErrors &&
+      chartAlert !== undefined &&
+      chartAlert !== NONEXISTENT_DATASET &&
+      datasource === PLACEHOLDER_DATASOURCE &&
+      datasetsStatus !== ResourceStatus.Error;
+
     return (
-      <ErrorContainer height={height}>
+      <ErrorContainer
+        height={height}
+        {...chartStatusAttrProps(
+          isDatasourceStillLoading
+            ? CHART_STATUS_ATTR_VALUES.LOADING
+            : CHART_STATUS_ATTR_VALUES.FAILED,
+        )}
+      >
         {queriesResponse?.map(item =>
           renderErrorMessage(item as ChartErrorType),
         )}
@@ -441,43 +476,53 @@ function Chart({
 
   if (chartStatus === 'stopped') {
     return (
-      <EmptyState
-        size="large"
-        title={chartAlert || t('Updating chart was stopped')}
-        description={
-          <span>
-            {t('Run a new query using the "Update chart" button or')}{' '}
-            <button
-              type="button"
-              onClick={onQuery}
-              css={css`
-                appearance: none;
-                border: none;
-                background: none;
-                padding: 0;
-                font: inherit;
-                cursor: pointer;
-                text-decoration: underline;
-              `}
-            >
-              {t('click here')}
-            </button>
-            .
-          </span>
-        }
-        image="chart.svg"
-      />
+      <div
+        {...chartStatusAttrProps(CHART_STATUS_ATTR_VALUES.STOPPED)}
+        css={{ width: '100%', height: '100%' }}
+      >
+        <EmptyState
+          size="large"
+          title={chartAlert || t('Updating chart was stopped')}
+          description={
+            <span>
+              {t('Run a new query using the "Update chart" button or')}{' '}
+              <button
+                type="button"
+                onClick={onQuery}
+                css={css`
+                  appearance: none;
+                  border: none;
+                  background: none;
+                  padding: 0;
+                  font: inherit;
+                  cursor: pointer;
+                  text-decoration: underline;
+                `}
+              >
+                {t('click here')}
+              </button>
+              .
+            </span>
+          }
+          image="chart.svg"
+        />
+      </div>
     );
   }
 
   if (errorMessage && ensureIsArray(queriesResponse).length === 0) {
     return (
-      <EmptyState
-        size="large"
-        title={t('Add required control values to preview chart')}
-        description={getChartRequiredFieldsMissingMessage(true)}
-        image="chart.svg"
-      />
+      <div
+        {...chartStatusAttrProps(CHART_STATUS_ATTR_VALUES.STOPPED)}
+        css={{ width: '100%', height: '100%' }}
+      >
+        <EmptyState
+          size="large"
+          title={t('Add required control values to preview chart')}
+          description={getChartRequiredFieldsMissingMessage(true)}
+          image="chart.svg"
+        />
+      </div>
     );
   }
   if (
@@ -488,34 +533,39 @@ function Chart({
     ensureIsArray(queriesResponse).length === 0
   ) {
     return (
-      <EmptyState
-        size="large"
-        title={t('Your chart is ready to go!')}
-        description={
-          <span>
-            {t(
-              'Click on "Create chart" button in the control panel on the left to preview a visualization or',
-            )}{' '}
-            <button
-              type="button"
-              onClick={onQuery}
-              css={css`
-                appearance: none;
-                border: none;
-                background: none;
-                padding: 0;
-                font: inherit;
-                cursor: pointer;
-                text-decoration: underline;
-              `}
-            >
-              {t('click here')}
-            </button>
-            .
-          </span>
-        }
-        image="chart.svg"
-      />
+      <div
+        {...chartStatusAttrProps(CHART_STATUS_ATTR_VALUES.STOPPED)}
+        css={{ width: '100%', height: '100%' }}
+      >
+        <EmptyState
+          size="large"
+          title={t('Your chart is ready to go!')}
+          description={
+            <span>
+              {t(
+                'Click on "Create chart" button in the control panel on the left to preview a visualization or',
+              )}{' '}
+              <button
+                type="button"
+                onClick={onQuery}
+                css={css`
+                  appearance: none;
+                  border: none;
+                  background: none;
+                  padding: 0;
+                  font: inherit;
+                  cursor: pointer;
+                  text-decoration: underline;
+                `}
+              >
+                {t('click here')}
+              </button>
+              .
+            </span>
+          }
+          image="chart.svg"
+        />
+      </div>
     );
   }
 
@@ -525,6 +575,7 @@ function Chart({
         data-ui-anchor="chart"
         className="chart-container"
         data-test="chart-container"
+        {...chartStatusAttrProps(getChartStatusAttrValue(chartStatus))}
         height={height}
         width={width}
       >

superset-frontend/src/utils/screenshotContract.ts

@@ -0,0 +1,95 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/**
+ * Single source of truth for the per-chart "screenshot readiness beacon"
+ * contract published on each chart's root DOM node.
+ *
+ * A backend screenshot/PDF pipeline (and a later frontend MutationObserver
+ * beacon) observes this attribute to decide when a chart has settled into a
+ * stable, paint-complete state. Every render branch of `Chart.tsx` must set
+ * `CHART_STATUS_ATTR` on its root DOM node to one of `CHART_STATUS_ATTR_VALUES`.
+ *
+ * This module intentionally holds only the attribute name, its allowed
+ * values, and the mapping from a Redux `chartStatus` to an attribute value.
+ * It must never reference CSS class names: `data-*` attributes have no
+ * styling hooks in this codebase, and this contract is not the place to
+ * introduce one.
+ */
+import type { ChartStatus } from 'src/explore/types';
+
+/** The DOM attribute name published on each chart's root node. */
+export const CHART_STATUS_ATTR = 'data-chart-status';
+
+/**
+ * Allowed values for `CHART_STATUS_ATTR`.
+ *
+ * - `pending`: no chart status yet, or a non-terminal/unrecognized status.
+ * - `loading`: a query is in flight, or data has arrived but the
+ *   visualization has not yet committed to the DOM.
+ * - `rendered`: the visualization has committed to the DOM (terminal).
+ * - `failed`: the query or render failed (terminal).
+ * - `stopped`: the query was stopped, or the chart is in a non-error
+ *   Explore-preview state (missing controls, ready-to-run) (terminal).
+ */
+export const CHART_STATUS_ATTR_VALUES = {
+  PENDING: 'pending',
+  LOADING: 'loading',
+  RENDERED: 'rendered',
+  FAILED: 'failed',
+  STOPPED: 'stopped',
+} as const;
+
+export type ChartStatusAttrValue =
+  (typeof CHART_STATUS_ATTR_VALUES)[keyof typeof CHART_STATUS_ATTR_VALUES];
+
+/**
+ * Maps a Redux `chartStatus` value to the `CHART_STATUS_ATTR` value for the
+ * shared success/loading/pending render branch of `Chart.tsx`.
+ *
+ * The `failed` and `stopped` chart statuses are handled by their own render
+ * branches (which may need additional local context, e.g. whether the
+ * datasource is still loading) and are included here only for completeness.
+ */
+export function getChartStatusAttrValue(
+  chartStatus: ChartStatus | '' | null | undefined,
+): ChartStatusAttrValue {
+  switch (chartStatus) {
+    case 'loading':
+      return CHART_STATUS_ATTR_VALUES.LOADING;
+    case 'success':
+      // Data has arrived but the visualization hasn't committed to the DOM
+      // yet. This is non-terminal; the terminal "rendered" value is
+      // published later, on render commit, via chartRenderingSucceeded.
+      return CHART_STATUS_ATTR_VALUES.LOADING;
+    case 'rendered':
+      return CHART_STATUS_ATTR_VALUES.RENDERED;
+    case 'failed':
+      return CHART_STATUS_ATTR_VALUES.FAILED;
+    case 'stopped':
+      return CHART_STATUS_ATTR_VALUES.STOPPED;
+    case null:
+    case undefined:
+    case '':
+    default:
+      // Fail-safe catch-all: pre-fetch, a failed-chart datasource-retry
+      // reset, or any other/future status not enumerated above.
+      return CHART_STATUS_ATTR_VALUES.PENDING;
+  }
+}

superset/commands/report/execute.py

@@ -788,11 +788,12 @@ class BaseReportState:
         try:
             imges = []
             for screenshot in screenshots:
-                imge = screenshot.get_screenshot(
+                capture_result = screenshot.get_screenshot(
                     user=user,
                     log_context=self._log_context,
                     report_execution_context=self._report_execution_context,
                 )
+                imge = capture_result.image
                 if imge is None:
                     raise ReportScheduleScreenshotFailedError(
                         "Screenshot failed; aborting to avoid sending a partial report"

superset/dashboards/excel_export/screenshot.py

@@ -89,7 +89,8 @@ def render_chart_image(
             window_size=window_size,
             thumb_size=window_size,
         )
-        return screenshot.get_screenshot(user=user)
+        capture_result = screenshot.get_screenshot(user=user)
+        return capture_result.image
     except SoftTimeLimitExceeded:
         # A soft timeout aborts the whole export; don't let the broad handler
         # below turn it into a ``None`` (a per-chart "could not render") result.

superset/tasks/cache.py

@@ -338,8 +338,8 @@ def _warmup_urls(
         for url in urls:
             try:
                 logger.info("Fetching %s", url)
-                screenshot = wd.get_screenshot(url, "grid-container", user=user)
-                if not screenshot:
+                capture_result = wd.get_screenshot(url, "grid-container", user=user)
+                if capture_result.image is None:
                     raise RuntimeError("No screenshot captured")
                 results["success"].append(url)
             except Exception:  # noqa: BLE001

superset/utils/screenshot_utils.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 import io
 import logging
 import time
+from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from celery import current_task
@@ -32,6 +33,23 @@ from superset.utils.report_execution import (
 
 logger = logging.getLogger(__name__)
 
+
+@dataclass(frozen=True)
+class ScreenshotCaptureResult:
+    """
+    Result of a screenshot capture attempt.
+
+    ``readiness`` is untyped (``object | None``) because the readiness-report
+    type does not exist yet; a later slice introduces it and narrows this
+    field's type. This dataclass intentionally has no ``__bool__`` override --
+    callers must test ``.image`` explicitly rather than the truthiness of the
+    result itself.
+    """
+
+    image: bytes | None
+    readiness: object | None = None
+
+
 # Time to wait after scrolling for content to settle and load (in milliseconds)
 SCROLL_SETTLE_TIMEOUT_MS = 1000
 

superset/utils/screenshots.py

@@ -34,6 +34,7 @@ from superset.exceptions import (
 from superset.extensions import event_logger
 from superset.utils.hashing import hash_from_dict
 from superset.utils.report_execution import ReportExecutionContext
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.urls import modify_url_query
 from superset.utils.webdriver import (
     ChartStandaloneMode,
@@ -216,16 +217,17 @@ class BaseScreenshot:
         window_size: WindowSize | None = None,
         log_context: str | None = None,
         report_execution_context: ReportExecutionContext | None = None,
-    ) -> bytes | None:
+    ) -> ScreenshotCaptureResult:
         driver = self.driver(window_size)
-        self.screenshot = driver.get_screenshot(
+        driver_result = driver.get_screenshot(
             self.url,
             self.element,
             user,
             log_context=log_context,
             report_execution_context=report_execution_context,
         )
-        return self.screenshot
+        self.screenshot = driver_result.image
+        return ScreenshotCaptureResult(image=self.screenshot, readiness=None)
 
     def get_cache_key(
         self,
@@ -326,11 +328,12 @@ class BaseScreenshot:
                     with event_logger.log_context(
                         f"screenshot.compute.{self.thumbnail_type}"
                     ):
-                        image = self.get_screenshot(
+                        capture_result = self.get_screenshot(
                             user=user,
                             window_size=window_size,
                             log_context=f"cache_key={cache_key}",
                         )
+                        image = capture_result.image
                 except Exception as ex:  # pylint: disable=broad-except
                     logger.warning(
                         "Failed at generating thumbnail for cache_key=%s: %s",

superset/utils/webdriver.py

@@ -37,6 +37,7 @@ from superset.utils.screenshot_utils import (
     FIND_CHART_HOLDER_STATES_JS,
     REPORT_CHART_HOLDERS_READY_JS,
     resolve_screenshot_task_budget_seconds,
+    ScreenshotCaptureResult,
     ScreenshotTaskBudgetExceededError,
     take_tiled_screenshot,
 )
@@ -155,7 +156,7 @@ class WebDriverProxy(ABC):
         user: User | None = None,
         log_context: str | None = None,
         report_execution_context: ReportExecutionContext | None = None,
-    ) -> bytes | None:
+    ) -> ScreenshotCaptureResult:
         """
         Run webdriver and return a screenshot
 
@@ -506,7 +507,7 @@ class WebDriverPlaywright(WebDriverProxy):
         user: User | None = None,
         log_context: str | None = None,
         report_execution_context: ReportExecutionContext | None = None,
-    ) -> bytes | None:
+    ) -> ScreenshotCaptureResult:
         screenshot_started_at = time.monotonic()
         if report_execution_context:
             log_context = report_execution_context.log_context
@@ -947,4 +948,4 @@ class WebDriverPlaywright(WebDriverProxy):
                 )
         finally:
             context.close()
-        return img
+        return ScreenshotCaptureResult(image=img, readiness=None)

tests/integration_tests/reports/commands/execute_dashboard_report_tests.py

@@ -27,6 +27,7 @@ from superset.commands.report.exceptions import ReportScheduleUnexpectedError
 from superset.commands.report.execute import AsyncExecuteReportScheduleCommand
 from superset.models.dashboard import Dashboard
 from superset.reports.models import ReportSourceFormat
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.urls import get_url_path
 from tests.integration_tests.fixtures.tabbed_dashboard import (
     tabbed_dashboard,  # noqa: F401
@@ -47,7 +48,9 @@ def test_report_for_dashboard_with_tabs(
     send_email_smtp_mock: MagicMock,
     tabbed_dashboard: Dashboard,  # noqa: F811
 ) -> None:
-    dashboard_screenshot_mock.get_screenshot.return_value = b"test-image"
+    dashboard_screenshot_mock.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=b"test-image", readiness=None
+    )
     current_app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"] = False
     with create_dashboard_report(
         dashboard=tabbed_dashboard,
@@ -90,7 +93,9 @@ def test_report_with_header_data(
     send_email_smtp_mock: MagicMock,
     tabbed_dashboard: Dashboard,  # noqa: F811
 ) -> None:
-    dashboard_screenshot_mock.get_screenshot.return_value = b"test-image"
+    dashboard_screenshot_mock.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=b"test-image", readiness=None
+    )
     current_app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"] = False
 
     with create_dashboard_report(
@@ -137,7 +142,9 @@ def test_report_schedule_stale_data_error_preserves_cause(
     we surface ReportScheduleUnexpectedError while preserving the original
     StaleDataError as the cause.
     """
-    dashboard_screenshot_mock.get_screenshot.return_value = b"test-image"
+    dashboard_screenshot_mock.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=b"test-image", readiness=None
+    )
     current_app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"] = False
 
     with create_dashboard_report(

tests/integration_tests/reports/commands_tests.py

@@ -89,6 +89,7 @@ from superset.tasks.types import ExecutorType
 from superset.utils import json
 from superset.utils.database import get_example_database
 from superset.utils.report_execution import ReportExecutionContext
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.webdriver import PlaywrightTimeout
 from tests.integration_tests.fixtures.birth_names_dashboard import (
     load_birth_names_dashboard_with_slices,  # noqa: F401
@@ -756,7 +757,9 @@ def test_email_chart_report_schedule_with_cc_bcc(
     ExecuteReport Command: Test chart email report schedule with screenshot and email cc, bcc options
     """  # noqa: E501
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -822,7 +825,9 @@ def test_email_chart_report_schedule(
     ExecuteReport Command: Test chart email report schedule with screenshot
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -868,7 +873,9 @@ def test_email_chart_report_schedule_single_log_per_execution(
     (the "trigger" row and the "result" row). This test asserts that one
     execution -- identified by its execution uuid -- yields exactly one log row.
     """
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -911,11 +918,11 @@ def test_email_chart_report_schedule_alpha_owner(
         user: User,
         log_context: Optional[str] = None,
         report_execution_context: ReportExecutionContext | None = None,
-    ) -> Optional[bytes]:
+    ) -> ScreenshotCaptureResult:
         nonlocal username
         username = user.username
 
-        return SCREENSHOT_FILE
+        return ScreenshotCaptureResult(image=SCREENSHOT_FILE, readiness=None)
 
     screenshot_mock.side_effect = _screenshot_side_effect
 
@@ -967,7 +974,9 @@ def test_email_chart_report_schedule_force_screenshot(
     reflect that.
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -1008,7 +1017,9 @@ def test_email_chart_alert_schedule(
     ExecuteReport Command: Test chart email alert schedule with screenshot
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -1045,7 +1056,9 @@ def test_email_chart_report_dry_run(
     ExecuteReport Command: Test chart email report schedule dry run
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     app.config["ALERT_REPORTS_NOTIFICATION_DRY_RUN"] = True
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -1169,7 +1182,9 @@ def test_email_chart_report_schedule_with_csv_no_query_context(
     ExecuteReport Command: Test chart email report schedule with CSV (no query context)
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     # setup csv mock
     response = Mock()
@@ -1332,7 +1347,9 @@ def test_email_dashboard_report_schedule(
     ExecuteReport Command: Test dashboard email report schedule
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         with patch(
@@ -1473,7 +1490,9 @@ def test_email_dashboard_report_schedule_force_screenshot(
     ExecuteReport Command: Test dashboard email report schedule
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -1512,7 +1531,9 @@ def test_slack_chart_report_schedule_converts_to_v2(
     while converting the recipients list to SlackV2.
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     channel_id = "slack_channel_id"
     get_channels_with_search_mock.return_value = [
         {
@@ -1573,7 +1594,9 @@ def test_slack_chart_report_schedule_converts_to_v2_channel_with_hash(
     the channel name includes the leading hash (supported in v1).
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     channel_id = "slack_channel_id"
     chart = db.session.query(Slice).first()
     report_schedule = create_report_notification(
@@ -1636,7 +1659,9 @@ def test_slack_chart_report_schedule_fails_to_converts_to_v2(
     ExecuteReport Command: Test converting a Slack report to v2 fails.
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     channel_id = "slack_channel_id"
     chart = db.session.query(Slice).first()
     report_schedule = create_report_notification(
@@ -1686,7 +1711,9 @@ def test_slack_chart_report_schedule_v2(
     ExecuteReport Command: Test chart slack report schedule using Slack v2.
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         with patch(
@@ -1725,7 +1752,9 @@ def test_slack_chart_report_schedule_with_errors(
     properly log something
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     slack_errors = [
         BotUserAccessError(),
@@ -2106,7 +2135,9 @@ def test_report_schedule_success_grace_end(
     ExecuteReport Command: Test report schedule on grace to noop
     """
 
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     # set current time to after the grace period
     current_time = create_alert_slack_chart_grace.last_eval_dttm + timedelta(
@@ -2144,7 +2175,9 @@ def test_alert_limit_is_applied(
     """
     ExecuteReport Command: Test that all alerts apply a SQL limit to stmts
     """
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with patch.object(
         create_alert_email_chart.database.db_engine_spec, "execute", return_value=None
@@ -2174,7 +2207,9 @@ def test_email_dashboard_report_fails(
     # setup screenshot mock
     from smtplib import SMTPException
 
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     email_mock.side_effect = SMTPException("Could not connect to SMTP XPTO")
 
     with pytest.raises(ReportScheduleSystemErrorsException):
@@ -2200,7 +2235,9 @@ def test_email_dashboard_report_fails_uncaught_exception(
     # setup screenshot mock
     from smtplib import SMTPException  # noqa: F401
 
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
     email_mock.side_effect = Exception("Uncaught exception")
     app.config["EMAIL_REPORTS_CTA"] = "Call to action"
 
@@ -2235,7 +2272,9 @@ def test_slack_chart_alert(
     ExecuteReport Command: Test chart slack alert
     """
     # setup screenshot mock
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     with freeze_time("2020-01-01T00:00:00Z"):
         AsyncExecuteReportScheduleCommand(
@@ -2319,7 +2358,9 @@ def test_slack_token_callable_chart_report(
     slack_token_mock = Mock(return_value="cool_code")
     with patch.dict("flask.current_app.config", {"SLACK_API_TOKEN": slack_token_mock}):
         # setup screenshot mock
-        screenshot_mock.return_value = SCREENSHOT_FILE
+        screenshot_mock.return_value = ScreenshotCaptureResult(
+            image=SCREENSHOT_FILE, readiness=None
+        )
 
         with freeze_time("2020-01-01T00:00:00Z"):
             AsyncExecuteReportScheduleCommand(
@@ -2598,7 +2639,9 @@ def test_readiness_timeout_retries_terminal_persistence_and_allows_next_schedule
 
     next_execution_id = str(uuid4())
     screenshot_mock.side_effect = None
-    screenshot_mock.return_value = SCREENSHOT_FILE
+    screenshot_mock.return_value = ScreenshotCaptureResult(
+        image=SCREENSHOT_FILE, readiness=None
+    )
 
     AsyncExecuteReportScheduleCommand(
         next_execution_id,
@@ -3230,7 +3273,9 @@ def test_success_after_retry_clears_retry_state(
 
     try:
         # Screenshot succeeds this time
-        screenshot_mock.return_value = SCREENSHOT_FILE
+        screenshot_mock.return_value = ScreenshotCaptureResult(
+            image=SCREENSHOT_FILE, readiness=None
+        )
 
         AsyncExecuteReportScheduleCommand(
             TEST_ID, report_schedule.id, scheduled_dttm

tests/unit_tests/commands/report/execute_test.py

@@ -70,6 +70,7 @@ from superset.utils.report_execution import (
     ReportExecutionContext,
     ReportExecutionDeadline,
 )
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.screenshots import ChartScreenshot
 from tests.integration_tests.conftest import with_feature_flags
 
@@ -1586,7 +1587,9 @@ def test_screenshot_width_calculation(
         # 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_screenshot.return_value = ScreenshotCaptureResult(
+            image=b"screenshot bytes", readiness=None
+        )
 
         # Mock get_executor to avoid database lookups
         with patch(

tests/unit_tests/dashboards/test_excel_export_screenshot.py

@@ -26,6 +26,7 @@ from superset.charts.data.dashboard_filter_context import DashboardFilterContext
 from superset.dashboards.excel_export import screenshot as screenshot_module
 from superset.dashboards.excel_export.screenshot import render_chart_image
 from superset.utils import json
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 
 MODULE = "superset.dashboards.excel_export.screenshot"
 
@@ -56,7 +57,9 @@ def test_render_chart_image_builds_url_with_slice_id_and_filters() -> None:
     )
 
     screenshot_instance = MagicMock()
-    screenshot_instance.get_screenshot.return_value = b"PNGBYTES"
+    screenshot_instance.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=b"PNGBYTES", readiness=None
+    )
 
     with (
         patch(
@@ -100,7 +103,9 @@ def test_render_chart_image_omits_extra_form_data_when_no_filters() -> None:
     chart = _chart()
     filter_context = DashboardFilterContext(extra_form_data={})
     screenshot_instance = MagicMock()
-    screenshot_instance.get_screenshot.return_value = b"PNG"
+    screenshot_instance.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=b"PNG", readiness=None
+    )
 
     with (
         patch(f"{MODULE}.get_dashboard_filter_context", return_value=filter_context),
@@ -118,7 +123,9 @@ def test_render_chart_image_omits_extra_form_data_when_no_filters() -> None:
 def test_render_chart_image_returns_none_when_screenshot_fails() -> None:
     chart = _chart()
     screenshot_instance = MagicMock()
-    screenshot_instance.get_screenshot.return_value = None
+    screenshot_instance.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=None, readiness=None
+    )
 
     with (
         patch(

tests/unit_tests/tasks/test_cache.py

@@ -17,6 +17,8 @@
 from typing import Any, Optional
 from unittest import mock
 
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
+
 
 def _fake_app(config: Optional[dict[str, Any]] = None) -> mock.MagicMock:
     """Build a stand-in for ``current_app`` with a controllable config dict."""
@@ -112,8 +114,9 @@ def test_cache_warmup_treats_none_screenshot_as_error(app_context: None) -> None
     urls: list[str] = ["http://localhost/dash/ok", "http://localhost/dash/none"]
     user: mock.MagicMock = mock.MagicMock()
 
-    def side_effect(url: str, _element: str, **kwargs: Any) -> Any:
-        return None if url.endswith("none") else b"PNG"
+    def side_effect(url: str, _element: str, **kwargs: Any) -> ScreenshotCaptureResult:
+        image = None if url.endswith("none") else b"PNG"
+        return ScreenshotCaptureResult(image=image, readiness=None)
 
     with (
         mock.patch(
@@ -142,10 +145,10 @@ def test_cache_warmup_collects_errors_and_destroys(app_context: None) -> None:
     urls: list[str] = ["http://localhost/dash/ok", "http://localhost/dash/boom"]
     user: mock.MagicMock = mock.MagicMock()
 
-    def side_effect(url: str, _element: str, **kwargs: Any) -> bytes | None:
+    def side_effect(url: str, _element: str, **kwargs: Any) -> ScreenshotCaptureResult:
         if url.endswith("boom"):
             raise Exception("screenshot failed")
-        return b"PNG"
+        return ScreenshotCaptureResult(image=b"PNG", readiness=None)
 
     with (
         mock.patch(

tests/unit_tests/utils/screenshot_test.py

@@ -23,6 +23,7 @@ import pytest
 from pytest_mock import MockerFixture
 
 from superset.utils.hashing import hash_from_dict
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.screenshots import (
     BaseScreenshot,
     ChartScreenshot,
@@ -69,12 +70,14 @@ def screenshot_obj():
 
 
 def test_get_screenshot(mocker: MockerFixture, screenshot_obj):
-    """Get screenshot should return a Bytes object"""
+    """Get screenshot should return a ScreenshotCaptureResult wrapping the bytes"""
     fake_bytes = b"fake_screenshot_data"
     driver = mocker.patch(BASE_SCREENSHOT_PATH + ".driver")
-    driver.return_value.get_screenshot.return_value = fake_bytes
+    driver.return_value.get_screenshot.return_value = ScreenshotCaptureResult(
+        image=fake_bytes, readiness=None
+    )
     screenshot_data = screenshot_obj.get_screenshot(mock_user)
-    assert screenshot_data == fake_bytes
+    assert screenshot_data.image == fake_bytes
 
 
 def test_get_cache_key(app_context, screenshot_obj):
@@ -111,7 +114,8 @@ class TestComputeAndCache:
             BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None
         )
         get_screenshot = mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
+            BASE_SCREENSHOT_PATH + ".get_screenshot",
+            return_value=ScreenshotCaptureResult(image=FAKE_PNG_BYTES, readiness=None),
         )
         resize_image = mocker.patch(
             BASE_SCREENSHOT_PATH + ".resize_image", return_value=FAKE_PNG_BYTES

tests/unit_tests/utils/test_screenshot_cache_fix.py

@@ -28,6 +28,7 @@ import pytest
 from pytest_mock import MockerFixture
 
 from superset.exceptions import LockAlreadyHeldException
+from superset.utils.screenshot_utils import ScreenshotCaptureResult
 from superset.utils.screenshots import (
     BaseScreenshot,
     ScreenshotCachePayload,
@@ -87,7 +88,8 @@ class TestCacheOnlyOnSuccess:
         mocker.patch(DISTRIBUTED_LOCK_PATH)
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         get_screenshot = mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
+            BASE_SCREENSHOT_PATH + ".get_screenshot",
+            return_value=ScreenshotCaptureResult(image=FAKE_PNG_BYTES, readiness=None),
         )
         # Mock resize_image to avoid PIL errors with fake image data
         mocker.patch(
@@ -171,7 +173,7 @@ class TestCacheOnlyOnSuccess:
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".get_screenshot",
-            return_value=b"",
+            return_value=ScreenshotCaptureResult(image=b"", readiness=None),
         )
         mock_logger = mocker.patch("superset.utils.screenshots.logger")
         BaseScreenshot.cache = MockCache()
@@ -200,7 +202,9 @@ class TestCacheOnlyOnSuccess:
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".get_screenshot",
-            return_value=b"this-is-not-a-real-image",
+            return_value=ScreenshotCaptureResult(
+                image=b"this-is-not-a-real-image", readiness=None
+            ),
         )
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".resize_image",
@@ -233,14 +237,16 @@ class TestCacheOnlyOnSuccess:
         mocker.patch(BASE_SCREENSHOT_PATH + ".get_from_cache_key", return_value=None)
         BaseScreenshot.cache = MockCache()
 
-        def check_cache_during_screenshot(*args: object, **kwargs: object) -> bytes:
+        def check_cache_during_screenshot(
+            *args: object, **kwargs: object
+        ) -> ScreenshotCaptureResult:
             cache_key = screenshot_obj.get_cache_key()
             cached_value = BaseScreenshot.cache.get(cache_key)
             assert cached_value is not None, (
                 "Cache should be set to COMPUTING before screenshot starts"
             )
             assert cached_value["status"] == "Computing"
-            return FAKE_PNG_BYTES
+            return ScreenshotCaptureResult(image=FAKE_PNG_BYTES, readiness=None)
 
         mocker.patch(
             BASE_SCREENSHOT_PATH + ".get_screenshot",
@@ -472,7 +478,8 @@ class TestIntegrationCacheBugFix:
         BaseScreenshot.cache.set(cache_key, stale_payload.to_dict())
 
         mocker.patch(
-            BASE_SCREENSHOT_PATH + ".get_screenshot", return_value=FAKE_PNG_BYTES
+            BASE_SCREENSHOT_PATH + ".get_screenshot",
+            return_value=ScreenshotCaptureResult(image=FAKE_PNG_BYTES, readiness=None),
         )
         # Mock resize to avoid PIL errors
         mocker.patch(

tests/unit_tests/utils/webdriver_test.py

@@ -149,7 +149,7 @@ class TestWebDriverPlaywrightFallback:
                 "http://example.com", "test-element", mock_user
             )
 
-        assert result == b"fake_screenshot"
+        assert result.image == b"fake_screenshot"
         mock_page.goto.assert_called_once_with(
             "http://example.com", wait_until="networkidle"
         )
@@ -201,7 +201,7 @@ class TestWebDriverPlaywrightFallback:
         # page.goto() timeout is caught and logged without aborting; execution
         # continues to the element waits, which succeed here, so a screenshot
         # is taken and returned (not None).
-        assert result is not None
+        assert result.image is not None
         mock_logger.exception.assert_called()
         exception_call = mock_logger.exception.call_args[0][0]
         assert "Web event %s not detected" in exception_call
@@ -535,7 +535,7 @@ class TestWebDriverPlaywrightErrorHandling:
                     "http://example.com", "dashboard", mock_user
                 )
 
-        assert result == b"fake_screenshot"
+        assert result.image == b"fake_screenshot"
         # chart_count (1) is well below the tiling threshold (20), so this is
         # the benign/expected case and must not be logged as a WARNING.
         mock_logger.debug.assert_any_call(
@@ -625,7 +625,7 @@ class TestWebDriverPlaywrightErrorHandling:
                     "http://example.com", "dashboard", mock_user
                 )
 
-        assert result == b"tiled_screenshot"
+        assert result.image == b"tiled_screenshot"
         mock_take_tiled.assert_called_once()
         mock_logger.warning.assert_any_call(
             "Could not determine dashboard height for element %s "
@@ -892,7 +892,7 @@ class TestWebDriverPlaywrightChartReadiness:
                 "http://example.com", "test-element", mock_user
             )
 
-        assert result == b"screenshot"
+        assert result.image == b"screenshot"
         # Readiness diagnostics are emitted before polling so a task killed by
         # an outer limit still leaves useful state in the logs.
         assert mock_page.evaluate.call_count == 2
@@ -990,7 +990,7 @@ class TestWebDriverPlaywrightChartReadiness:
 
         predicate = mock_page.wait_for_function.call_args.args[0]
         assert "holders.length > 0" not in predicate
-        assert result == mock_page.screenshot.return_value
+        assert result.image == mock_page.screenshot.return_value
 
     @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True)
     @patch("superset.utils.webdriver._browser_manager")
@@ -1467,7 +1467,7 @@ class TestWebDriverPlaywrightAnimationWaitOrder:
                 "http://example.com", "test-element", mock_user
             )
 
-        assert result == b"screenshot"
+        assert result.image == b"screenshot"
         mock_take_tiled.assert_not_called()
         mock_page.set_viewport_size.assert_not_called()
 
@@ -1500,7 +1500,7 @@ class TestWebDriverPlaywrightAnimationWaitOrder:
                 "http://example.com", "standalone", mock_user
             )
 
-        assert result == b"tiled_screenshot"
+        assert result.image == b"tiled_screenshot"
         mock_take_tiled.assert_called_once_with(
             mock_page,
             "standalone",