Migrate to test-state attribute to capture e2e script state (streamli…

…t#9131)

## Describe your changes

The e2e playwright tests currently rely on the status widget to check
for the running state. This PR changes that to use the internal
`data-teststate` on `stApp` instead which is more accurate. Also, we
might soon change the behaviour of the status widget.

## Testing Plan

- This only involves changes to our e2e tests.

---

**Contribution License Agreement**

By submitting this pull request you agree that all contributions to this
project are made under the Apache 2.0 license.

e2e_playwright/app_hotkeys.py

@@ -12,6 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 import streamlit as st
 
+if "counter" not in st.session_state:
+    st.session_state.counter = 0
+
+st.session_state.counter += 1
+
 st.text_input("text_input")
+
+st.write("Script runs:", st.session_state.counter)

e2e_playwright/app_hotkeys_test.py

@@ -14,6 +14,11 @@
 
 from playwright.sync_api import Page, expect
 
+from e2e_playwright.conftest import wait_for_app_run
+from e2e_playwright.shared.app_utils import (
+    expect_prefixed_markdown,
+)
+
 
 def test_shows_clear_cache_dialog_when_c_is_pressed(app: Page):
     app.keyboard.type("c")
@@ -40,12 +45,15 @@ def test_does_not_clear_cache_dialog_when_c_is_pressed_inside_text_input(app: Pa
 
 
 def test_reruns_when_r_is_pressed(app: Page):
+    expect_prefixed_markdown(app, "Script runs:", "1", exact_match=False)
     app.keyboard.type("r")
-    expect(app.get_by_test_id("stStatusWidget")).to_be_visible()
+    wait_for_app_run(app)
+    expect_prefixed_markdown(app, "Script runs:", "2", exact_match=False)
 
 
-def test_does_not_clear_cache_dialog_when_r_is_pressed_inside_text_input(
+def test_does_not_rerun_when_r_is_pressed_inside_text_input(
     app: Page,
 ):
-    app.get_by_test_id("stTextInput").press("r")
-    expect(app.get_by_test_id("stStatusWidget")).not_to_be_visible()
+    app.get_by_test_id("stTextInput").locator("input").press("r")
+    wait_for_app_run(app)
+    expect_prefixed_markdown(app, "Script runs:", "1", exact_match=False)

e2e_playwright/conftest.py

@@ -612,9 +612,19 @@ def wait_for_app_run(page: Page, wait_delay: int = 100):
     """Wait for the given page to finish running."""
     # Add a little timeout to wait for eventual debounce timeouts used in some widgets.
     page.wait_for_timeout(155)
-
+    # Make sure that the websocket connection is established.
+    page.wait_for_selector(
+        "[data-testid='stApp'][data-test-connection-state='CONNECTED']",
+        timeout=20000,
+        state="attached",
+    )
+    # Wait until we know the script has started. We determine this by checking
+    # whether the app is in notRunning state. (The data-test-connection-state attribute
+    # goes through the sequence "initial" -> "running" -> "notRunning").
     page.wait_for_selector(
-        "[data-testid='stStatusWidget']", timeout=20000, state="detached"
+        "[data-testid='stApp'][data-test-script-state='notRunning']",
+        timeout=20000,
+        state="attached",
     )
 
     if wait_delay > 0:

e2e_playwright/shared/app_utils.py

@@ -16,7 +16,7 @@
 
 import platform
 import re
-from typing import Pattern
+from typing import Literal, Pattern
 
 from playwright.sync_api import Locator, Page, expect
 
@@ -410,3 +410,31 @@ def expect_help_tooltip(
         position={"x": 0, "y": 0}, no_wait_after=True, force=True
     )
     expect(tooltip_content).not_to_be_attached()
+
+
+def expect_script_state(
+    page: Page,
+    state: Literal[
+        "initial",
+        "running",
+        "notRunning",
+        "rerunRequested",
+        "stopRequested",
+        "compilationError",
+    ],
+) -> None:
+    """Expect the app to be in a specific script state.
+
+    Parameters
+    ----------
+    page : Page
+        The page to search for the script state.
+
+    state :
+        The expected script state.
+    """
+    page.wait_for_selector(
+        f"[data-testid='stApp'][data-test-script-state='{state}']",
+        timeout=10000,
+        state="attached",
+    )

frontend/app/src/App.test.tsx

@@ -343,6 +343,48 @@ describe("App", () => {
     expect(getMockConnectionManager().disconnect).toHaveBeenCalled()
   })
 
+  it("correctly sets the data-test-script-state attribute", async () => {
+    renderApp(getProps())
+
+    expect(screen.getByTestId("stApp")).toHaveAttribute(
+      "data-test-script-state",
+      "initial"
+    )
+
+    sendForwardMessage("newSession", NEW_SESSION_JSON)
+
+    sendForwardMessage("sessionStatusChanged", {
+      runOnSave: false,
+      scriptIsRunning: true,
+    })
+
+    await waitFor(() => {
+      expect(screen.getByTestId("stApp")).toHaveAttribute(
+        "data-test-script-state",
+        ScriptRunState.RUNNING
+      )
+    })
+
+    sendForwardMessage("sessionStatusChanged", {
+      runOnSave: false,
+      scriptIsRunning: false,
+    })
+
+    expect(screen.getByTestId("stApp")).toHaveAttribute(
+      "data-test-script-state",
+      ScriptRunState.NOT_RUNNING
+    )
+
+    sendForwardMessage("sessionEvent", {
+      type: "scriptCompilationException",
+    })
+
+    expect(screen.getByTestId("stApp")).toHaveAttribute(
+      "data-test-script-state",
+      ScriptRunState.COMPILATION_ERROR
+    )
+  })
+
   describe("streamlit server version changes", () => {
     let prevWindowLocation: Location
 
@@ -1976,6 +2018,39 @@ describe("App", () => {
       })
     })
 
+    it("Correctly sets the data-test-connection-state attribute", () => {
+      renderApp(getProps())
+
+      const connectionManager = getMockConnectionManager(false)
+
+      expect(screen.getByTestId("stApp")).toHaveAttribute(
+        "data-test-connection-state",
+        ConnectionState.INITIAL
+      )
+
+      act(() =>
+        // @ts-expect-error - connectionManager.props is private
+        connectionManager.props.connectionStateChanged(
+          ConnectionState.CONNECTED
+        )
+      )
+      expect(screen.getByTestId("stApp")).toHaveAttribute(
+        "data-test-connection-state",
+        ConnectionState.CONNECTED
+      )
+
+      act(() =>
+        // @ts-expect-error - connectionManager.props is private
+        connectionManager.props.connectionStateChanged(
+          ConnectionState.PINGING_SERVER
+        )
+      )
+      expect(screen.getByTestId("stApp")).toHaveAttribute(
+        "data-test-connection-state",
+        ConnectionState.PINGING_SERVER
+      )
+    })
+
     it("Sets attemptingToReconnect to false if DISCONNECTED_FOREVER", () => {
       renderApp(getProps())
 

frontend/app/src/App.tsx

@@ -1908,11 +1908,12 @@ export class App extends PureComponent<Props, State> {
             <StyledApp
               className={outerDivClass}
               data-testid="stApp"
-              data-teststate={
+              data-test-script-state={
                 scriptRunId == INITIAL_SCRIPT_RUN_ID
                   ? "initial"
                   : scriptRunState
               }
+              data-test-connection-state={connectionState}
             >
               {/* The tabindex below is required for testing. */}
               <Header>

frontend/cypress/support/commands.js

@@ -168,9 +168,9 @@ Cypress.Commands.add("loadApp", (appUrl, timeout) => {
 
 Cypress.Commands.add("waitForScriptFinish", (timeout = 20000) => {
   // Wait until we know the script has started. We determine this by checking
-  // whether the app is in notRunning state. (The data-teststate attribute goes
-  // through the sequence "initial" -> "running" -> "notRunning")
-  cy.get("[data-testid='stApp'][data-teststate='notRunning']", {
+  // whether the app is in notRunning state. (The data-test-connection-state attribute
+  // goes through the sequence "initial" -> "running" -> "notRunning")
+  cy.get("[data-testid='stApp'][data-test-script-state='notRunning']", {
     timeout,
   }).should("exist")
 })

frontend/lib/src/ScriptRunState.ts

@@ -15,7 +15,7 @@
  */
 
 export enum ScriptRunState {
-  // TODO: Add INITIAL state here and clean up data-teststate in App.tsx.
+  // TODO: Add INITIAL state here and clean up data-test-script-state in App.tsx.
   // But before we do this, we need to make sure Snowflake hosts that use this
   // state will not break. And that's a bigger project...
   //INITIAL = "initial",