e2e_playwright/st_pydeck_chart_select_test.py - develop - vshmidt/streamlit

e2e_playwright/st_pydeck_chart_select_test.py
487 строк16 KB
Bob Nisco
[fix] Assert Snapshot pixel_threshold must be meaningful (#13896)
11 фев 2026, 02:10
11 фев 2026, 02:10d55cd7d
# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026)
#
# Licensed 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.


from __future__ import annotations

from typing import Literal

import pytest
from playwright.sync_api import Locator, Page, Position, expect

from e2e_playwright.conftest import (
    ImageCompareFunction,
    wait_for_app_run,
)
from e2e_playwright.shared.app_utils import (
    click_button,
    click_form_button,
    expect_prefixed_markdown,
)
from e2e_playwright.shared.pydeck_utils import (
    click_point,
    get_click_handling_div,
    wait_for_chart,
)

# The pydeck tests are a lot flakier than need be so increase the pixel threshold
PIXEL_THRESHOLD = 0.1

# Common selection values
EMPTY_SELECTION = "{'selection': {'indices': {}, 'objects': {}}}"
FIRST_POINT_SELECTION = (
    "{'selection': {'indices': {'MyHexLayer': [0]}, "
    "'objects': {'MyHexLayer': [{'count': 10, 'hex': '88283082b9fffff'}]}}}"
)
SECOND_POINT_SELECTION = (
    "{'selection': {'indices': {'MyHexLayer': [2]}, "
    "'objects': {'MyHexLayer': [{'count': 100, 'hex': '88283082a9fffff'}]}}}"
)
MULTI_SELECTION = (
    "{'selection': {'indices': {'MyHexLayer': [0, 2]}, "
    "'objects': {'MyHexLayer': [{'count': 10, 'hex': '88283082b9fffff'}, "
    "{'count': 100, 'hex': '88283082a9fffff'}]}}}"
)


FIRST_POINT_COORDS: Position = {"x": 344.0, "y": 201.0}
SECOND_POINT_COORDS: Position = {"x": 417.0, "y": 229.0}
DESELECT_COORDS: Position = {"x": 0.0, "y": 0.0}
SCATTERPLOT_POINT_COORDS: Position = {"x": 279.0, "y": 331.0}
# Coordinates for the form test (slightly different from FIRST_POINT_COORDS)
FORM_POINT_COORDS: Position = {"x": 326.0, "y": 208.0}

# Standard wait delay for app runs after interactions
STANDARD_WAIT_DELAY = 5000


def _select_chart_type(app: Page, chart_type: str):
    """Select the chart type to display."""
    button_text_map = {
        "basic": "Basic Chart",
        "callback": "With Callback",
        "form": "In Form",
        "fragment": "In Fragment",
        "scatterplot": "Scatterplot",
        "dynamic": "Dynamic Data",
    }

    button_text = button_text_map.get(chart_type, "Basic Chart")
    click_button(app, button_text)


def _set_selection_mode(app: Page, mode: Literal["single-object", "multi-object"]):
    """Set the selection mode for the PyDeck chart."""
    app.get_by_test_id("stSelectbox").nth(0).locator("input").click()
    selection_dropdown = app.locator('[data-baseweb="popover"]').first
    selection_dropdown.locator("li").nth(1 if mode == "multi-object" else 0).click()

    # click elsewhere to close the dropdown
    app.get_by_test_id("stApp").click(position={"x": 0, "y": 0})

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)


def _click_point_and_verify_selection(
    app: Page,
    click_handling_div: Locator,
    coords: Position,
    expected_selection: str,
    markdown_prefix: str = "managed_map selection:",
    markdown_prefix_session_state: str | None = "session_state.managed_map:",
    wait_delay: int = STANDARD_WAIT_DELAY,
    max_attempts: int = 1,
):
    """Helper function to click on a point and verify the selection."""
    for attempt in range(max_attempts):
        click_point(click_handling_div, coords)

        wait_for_app_run(app, wait_delay=wait_delay)

        try:
            if markdown_prefix:
                expect_prefixed_markdown(
                    app,
                    markdown_prefix,
                    expected_selection,
                )

            if markdown_prefix_session_state:
                expect_prefixed_markdown(
                    app,
                    markdown_prefix_session_state,
                    expected_selection,
                )
            return
        except AssertionError:
            if attempt == max_attempts - 1:
                raise


# A note on browser testing strategy. We are only testing on Chromium because:
#   - Firefox seems to be failing but can't reproduce locally and video from CI
#     produces an empty element for PyDeck
#   - Webkit is too slow on CI, leading to flaky tests
#
# Getting coverage in Chromium is sufficient for now since the underlying logic
# is the same and we are not testing for browser-specific behavior.
#
# Additionally, even the deck.gl docs caution against visual tests since it
# renders to a canvas element:
# https://github.com/visgl/deck.gl/blob/master/docs/developer-guide/testing.md


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_multiselect_interactions_and_return_values(app: Page):
    """
    Test single selection, multi selection, and deselection all function
    properly and return the expected values in both session_state and as a
    return of st.pydeck.
    """
    _select_chart_type(app, "basic")
    _set_selection_mode(app, "multi-object")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)
    markdown_prefix_session_state = "session_state.managed_map:"
    markdown_prefix = "managed_map selection:"

    # Assert we haven't yet written anything out for the debugging state
    expect_prefixed_markdown(
        app,
        markdown_prefix,
        EMPTY_SELECTION,
    )
    expect_prefixed_markdown(
        app,
        markdown_prefix_session_state,
        EMPTY_SELECTION,
    )

    # Click on the hex that has count: 10
    _click_point_and_verify_selection(
        app, click_handling_div, FIRST_POINT_COORDS, FIRST_POINT_SELECTION
    )

    # Multiselect and click the hex that has count: 100
    _click_point_and_verify_selection(
        app, click_handling_div, SECOND_POINT_COORDS, MULTI_SELECTION
    )

    # Deselect everything by clicking away from an object in a layer
    _click_point_and_verify_selection(
        app, click_handling_div, DESELECT_COORDS, EMPTY_SELECTION
    )


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_single_select_interactions_and_return_values(
    app: Page, assert_snapshot: ImageCompareFunction
):
    """
    Test single selection and deselection all function properly and return the
    expected values in both session_state and as a return of st.pydeck.
    """
    _select_chart_type(app, "basic")
    _set_selection_mode(app, "single-object")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    # Click on the hex that has count: 10
    _click_point_and_verify_selection(
        app, click_handling_div, FIRST_POINT_COORDS, FIRST_POINT_SELECTION
    )

    # Click the hex that has count: 100
    _click_point_and_verify_selection(
        app, click_handling_div, SECOND_POINT_COORDS, SECOND_POINT_SELECTION
    )

    # Click on the hex that has count: 10 again
    _click_point_and_verify_selection(
        app, click_handling_div, FIRST_POINT_COORDS, FIRST_POINT_SELECTION
    )

    # Deselect everything by clicking away from an object in a layer
    _click_point_and_verify_selection(
        app, click_handling_div, DESELECT_COORDS, EMPTY_SELECTION
    )

    # Scatterplot checks
    _select_chart_type(app, "scatterplot")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    # Click on the scatterplot point with the biggest size
    _click_point_and_verify_selection(
        app,
        click_handling_div,
        SCATTERPLOT_POINT_COORDS,
        "'indices': {'cities': [4]}",
        markdown_prefix="scatterplot_map selection:",
        markdown_prefix_session_state="session_state.scatterplot_map:",
        # Scatterplot selection can be weird in CI, allow it to retry a few
        # times to avoid flakes.
        max_attempts=3,
    )

    # Assert that we have deselected everything
    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_select-scatterplot-single-selection",
        pixel_threshold=PIXEL_THRESHOLD,
    )


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_multiselect_has_consistent_visuals(
    app: Page, assert_snapshot: ImageCompareFunction
):
    """
    Test that no selection, single selection, multi selection, and deselection
    all look visually correct.
    """
    _select_chart_type(app, "basic")
    _set_selection_mode(app, "multi-object")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_select-no-selections",
        pixel_threshold=PIXEL_THRESHOLD,
    )

    # Click on the hex that has count: 10
    click_point(click_handling_div, FIRST_POINT_COORDS)

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_select-single-selection",
        pixel_threshold=PIXEL_THRESHOLD,
    )

    # Multiselect and click the hex that has count: 100
    click_point(click_handling_div, SECOND_POINT_COORDS)

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_select-multi-selection",
        pixel_threshold=PIXEL_THRESHOLD,
    )

    # Deselect everything by clicking away from an object in a layer
    click_point(click_handling_div, DESELECT_COORDS)

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    # Assert that we have deselected everything
    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_select-deselected",
        pixel_threshold=PIXEL_THRESHOLD,
    )


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_selection_state_remains_after_unmounting(
    app: Page, assert_snapshot: ImageCompareFunction
):
    """
    Test that no selection, single selection, multi selection, and deselection
    all look visually correct.
    """
    _select_chart_type(app, "basic")
    _set_selection_mode(app, "multi-object")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    # Click on the hex that has count: 10
    click_point(click_handling_div, FIRST_POINT_COORDS)

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    # Multiselect and click the hex that has count: 100
    click_point(click_handling_div, SECOND_POINT_COORDS)

    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    click_button(app, "Create some elements to unmount component")
    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)

    wait_for_chart(app)

    assert_snapshot(
        click_handling_div,
        name="st_pydeck_chart_selection_state_remains_after_unmounting",
        pixel_threshold=PIXEL_THRESHOLD,
    )


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_selection_callback(app: Page):
    """Test the callback functionality of a PyDeck chart."""
    _select_chart_type(app, "callback")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    markdown_prefix = "PyDeck selection callback:"

    # Assert we haven't yet written anything out for the debugging state
    expect(app.get_by_text(markdown_prefix)).to_have_count(0)

    # Click on the hex that has count: 10
    _click_point_and_verify_selection(
        app,
        click_handling_div,
        FIRST_POINT_COORDS,
        FIRST_POINT_SELECTION,
        markdown_prefix,
        markdown_prefix_session_state=None,
    )


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_selection_in_form(app: Page):
    """Test the selection functionality of a PyDeck chart within a form."""
    _select_chart_type(app, "form")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    # Assert we haven't yet written anything out for the debugging state
    markdown_prefix = "PyDeck-in-form selection:"
    markdown_prefix_session_state = "PyDeck-in-form selection in session state:"

    # Click on the hex that has count: 10
    click_point(click_handling_div, FORM_POINT_COORDS)

    wait_for_app_run(app)

    # Nothing should be shown yet because we did not submit the form
    expect_prefixed_markdown(
        app,
        markdown_prefix,
        EMPTY_SELECTION,
    )
    expect_prefixed_markdown(
        app,
        markdown_prefix_session_state,
        EMPTY_SELECTION,
    )

    # submit the form. The selection uses a debounce of 200ms; if we click too
    # early, the state is not updated correctly and we submit the old,
    # unselected values
    app.wait_for_timeout(210)

    click_form_button(app, "Submit")

    expect_prefixed_markdown(app, markdown_prefix, FIRST_POINT_SELECTION)
    expect_prefixed_markdown(app, markdown_prefix_session_state, FIRST_POINT_SELECTION)
    expect(app.get_by_test_id("stForm")).not_to_contain_text("Error")


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_selection_in_fragment(app: Page):
    """Test the selection functionality of a PyDeck chart within a fragment."""
    _select_chart_type(app, "fragment")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    # Check that the main script has run twice (the initial run, and the run
    # after selecting the Fragment type)
    expect(app.get_by_text("Runs: 2")).to_be_visible()

    # Assert we haven't yet written anything out for the debugging state
    markdown_prefix = "PyDeck-in-fragment selection:"

    # Nothing should be shown yet because we did do anything yet
    expect_prefixed_markdown(
        app,
        markdown_prefix,
        EMPTY_SELECTION,
    )

    # Click on the hex that has count: 10
    _click_point_and_verify_selection(
        app,
        click_handling_div,
        FIRST_POINT_COORDS,
        FIRST_POINT_SELECTION,
        markdown_prefix,
        markdown_prefix_session_state=None,
    )

    # Check that the main script has not re-run any additional times.
    expect(app.get_by_text("Runs: 2")).to_be_visible()


@pytest.mark.only_browser("chromium")
@pytest.mark.flaky(reruns=4)
def test_pydeck_chart_selection_preserved_when_data_updates(app: Page):
    """
    Test that selection state is preserved when data updates and a key is provided.
    This tests the key_as_main_identity={"selection_mode"} behavior.
    """
    _select_chart_type(app, "dynamic")
    wait_for_chart(app)

    click_handling_div = get_click_handling_div(app, nth=0)

    markdown_prefix = "dynamic_chart selection:"
    markdown_prefix_session_state = "session_state.dynamic_chart:"

    # Verify initial empty state
    expect_prefixed_markdown(app, markdown_prefix, EMPTY_SELECTION)
    expect_prefixed_markdown(app, markdown_prefix_session_state, EMPTY_SELECTION)

    # Select the first point
    _click_point_and_verify_selection(
        app,
        click_handling_div,
        FIRST_POINT_COORDS,
        FIRST_POINT_SELECTION,
        markdown_prefix,
        markdown_prefix_session_state,
    )

    # Click "Update Data" button to change the data (but keep same key)
    click_button(app, "Update Data")
    wait_for_app_run(app, wait_delay=STANDARD_WAIT_DELAY)
    wait_for_chart(app)

    # Verify data version changed
    expect(app.get_by_text("Data version: 2")).to_be_visible()

    # Selection should still be preserved because key is provided
    expect_prefixed_markdown(app, markdown_prefix, FIRST_POINT_SELECTION)
    expect_prefixed_markdown(app, markdown_prefix_session_state, FIRST_POINT_SELECTION)