e2e_playwright/shared/pydeck_utils.py - develop - vshmidt/streamlit

e2e_playwright/shared/pydeck_utils.py
95 строк3 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 TYPE_CHECKING, Literal

import pandas as pd
import pydeck as pdk
from playwright.sync_api import Locator, Page, Position, expect

import streamlit as st

if TYPE_CHECKING:
    from streamlit.elements.deck_gl_json_chart import PydeckState
    from streamlit.runtime.state.common import WidgetCallback

H3_HEX_DATA = [
    {"hex": "88283082b9fffff", "count": 10},
    {"hex": "88283082d7fffff", "count": 50},
    {"hex": "88283082a9fffff", "count": 100},
]
df = pd.DataFrame(H3_HEX_DATA)


def get_pydeck_chart(
    key: str,
    selection_mode: Literal["single-object", "multi-object"],
    on_select: WidgetCallback | None = None,
) -> PydeckState:
    return st.pydeck_chart(
        pdk.Deck(
            initial_view_state=pdk.ViewState(
                latitude=37.7749295,
                longitude=-122.4194155,
                zoom=11,
                bearing=0,
                pitch=30,
            ),
            layers=[
                pdk.Layer(
                    "H3HexagonLayer",
                    df,
                    id="MyHexLayer",
                    stroked=True,
                    filled=True,
                    get_hexagon="hex",
                    line_width_min_pixels=2,
                    get_fill_color="[120, count > 50 ? 255 : 0, 255]",
                ),
            ],
        ),
        width="stretch",
        key=key,
        on_select=on_select or "rerun",
        selection_mode=selection_mode,
    )


def wait_for_chart(app: Page) -> None:
    # The pydeck chart takes a while to load so check that
    # it gets attached with an increased timeout.
    pydeck_charts = app.get_by_test_id("stDeckGlJsonChart")
    expect(pydeck_charts).to_have_count(1, timeout=10000)

    # The map assets can take more time to load and render especially in CI due
    # to the underlying hardware. Add an extra timeout to naively prevent
    # flakiness.
    app.wait_for_timeout(10000)


def get_click_handling_div(app: Page, nth: int) -> Locator:
    # Find canvas with class name "mapboxgl-canvas"
    expect(app.locator(".mapboxgl-canvas").nth(nth)).to_be_visible()
    click_handling_div = app.locator("#view-default-view").nth(nth)
    click_handling_div.scroll_into_view_if_needed()
    return click_handling_div


def click_point(click_handling_div: Locator, coords: Position) -> None:
    """Helper function to click on a point."""
    # Use force=True since it seems like another div sometimes intercepts events
    # in CI, causing the click to fail
    click_handling_div.click(position=coords, force=True)