frontend/lib/src/hostComm/HostCommunicationManager.test.tsx - develop - vshmidt/streamlit

frontend/lib/src/hostComm/HostCommunicationManager.test.tsx
787 строк22 KB
Lukas Masuch
Remove deprecated proto fields and unused theme creator code (#13842)
10 фев 2026, 01:13
10 фев 2026, 01:138722152
/**
 * 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.
 */

import "../../../utils/src/polyfills/index"

import { MockInstance } from "vitest"

import HostCommunicationManager, {
  HOST_COMM_VERSION,
} from "~lib/hostComm/HostCommunicationManager"

// Mocking "message" event listeners on the window;
// returns function to establish a listener
// eslint-disable-next-line @typescript-eslint/no-explicit-any -- TODO: Replace 'any' with a more specific type.
function mockEventListeners(): (type: string, event: any) => void {
  const listeners: { [name: string]: ((event: Event) => void)[] } = {}

  // eslint-disable-next-line @typescript-eslint/no-explicit-any -- TODO: Replace 'any' with a more specific type.
  window.addEventListener = vi.fn((event: string, cb: any) => {
    listeners[event] = listeners[event] || []
    listeners[event].push(cb)
  })

  return (type: string, event: Event): void =>
    listeners[type].forEach(cb => cb(event))
}

describe("HostCommunicationManager messaging", () => {
  let hostCommunicationMgr: HostCommunicationManager

  let dispatchEvent: (type: string, event: Event) => void
  let originalHash: string

  let setAllowedOriginsFunc: MockInstance
  let openCommFunc: MockInstance
  let sendMessageToHostFunc: MockInstance

  beforeEach(() => {
    hostCommunicationMgr = new HostCommunicationManager({
      streamlitExecutionStartedAt: 100,
      themeChanged: vi.fn(),
      sendRerunBackMsg: vi.fn(),
      pageChanged: vi.fn(),
      closeModal: vi.fn(),
      stopScript: vi.fn(),
      rerunScript: vi.fn(),
      clearCache: vi.fn(),
      sendAppHeartbeat: vi.fn(),
      setInputsDisabled: vi.fn(),
      isOwnerChanged: vi.fn(),
      fileUploadClientConfigChanged: vi.fn(),
      hostMenuItemsChanged: vi.fn(),
      hostToolbarItemsChanged: vi.fn(),
      hostHideSidebarNavChanged: vi.fn(),
      sidebarChevronDownshiftChanged: vi.fn(),
      pageLinkBaseUrlChanged: vi.fn(),
      queryParamsChanged: vi.fn(),
      deployedAppMetadataChanged: vi.fn(),
      restartWebsocketConnection: vi.fn(),
      terminateWebsocketConnection: vi.fn(),
    })

    originalHash = window.location.hash
    dispatchEvent = mockEventListeners()

    setAllowedOriginsFunc = vi.spyOn(hostCommunicationMgr, "setAllowedOrigins")
    openCommFunc = vi.spyOn(hostCommunicationMgr, "openHostCommunication")
    sendMessageToHostFunc = vi.spyOn(hostCommunicationMgr, "sendMessageToHost")

    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["https://devel.streamlit.test"],
      useExternalAuthToken: false,
    })
  })

  afterEach(() => {
    window.location.hash = originalHash
  })

  it("sets allowedOrigins properly & opens HostCommunication", () => {
    expect(setAllowedOriginsFunc).toHaveBeenCalledWith({
      allowedOrigins: ["https://devel.streamlit.test"],
      useExternalAuthToken: false,
    })
    // @ts-expect-error
    expect(hostCommunicationMgr.allowedOrigins).toEqual([
      "https://devel.streamlit.test",
    ])
    expect(openCommFunc).toHaveBeenCalled()
  })

  it("host should receive a GUEST_READY message", () => {
    expect(sendMessageToHostFunc).toHaveBeenCalled()

    const guestReadyMessage = sendMessageToHostFunc.mock.calls[0][0]
    expect(guestReadyMessage).toHaveProperty("type", "GUEST_READY")
    expect(guestReadyMessage).toHaveProperty("streamlitExecutionStartedAt")
    expect(guestReadyMessage).toHaveProperty(
      "guestReadyAt",
      expect.any(Number)
    )
  })

  it("can process a received CLOSE_MODAL message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "CLOSE_MODAL",
        },
        origin: "https://devel.streamlit.test",
      })
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.closeModal).toHaveBeenCalled()
  })

  it("can process a received STOP_SCRIPT message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "STOP_SCRIPT",
        },
        origin: "https://devel.streamlit.test",
      })
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.stopScript).toHaveBeenCalled()
  })

  it("can process a received RERUN_SCRIPT message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "RERUN_SCRIPT",
        },
        origin: "https://devel.streamlit.test",
      })
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.rerunScript).toHaveBeenCalled()
  })

  it("can process a received CLEAR_CACHE message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "CLEAR_CACHE",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.clearCache).toHaveBeenCalled()
  })

  it("can process a received REQUEST_PAGE_CHANGE message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "REQUEST_PAGE_CHANGE",
          pageScriptHash: "hash1",
        },
        origin: "https://devel.streamlit.test",
      })
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.pageChanged).toHaveBeenCalledWith(
      "hash1"
    )
  })

  it("can process a received SEND_APP_HEARTBEAT message without ackTimeoutMilliseconds", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SEND_APP_HEARTBEAT",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.sendAppHeartbeat).toHaveBeenCalledWith(0)
  })

  it("can process a received SEND_APP_HEARTBEAT message with ackTimeoutMilliseconds", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SEND_APP_HEARTBEAT",
          ackTimeoutMilliseconds: 59000,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.sendAppHeartbeat).toHaveBeenCalledWith(
      59000
    )
  })

  it("treats negative ackTimeoutMilliseconds as 0", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SEND_APP_HEARTBEAT",
          ackTimeoutMilliseconds: -1000,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.sendAppHeartbeat).toHaveBeenCalledWith(0)
  })

  it("can process a received SET_INPUTS_DISABLED message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_INPUTS_DISABLED",
          disabled: true,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.setInputsDisabled).toHaveBeenCalledWith(
      true
    )
  })

  it("should respond to SET_IS_OWNER message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_IS_OWNER",
          isOwner: true,
        },
        origin: "https://devel.streamlit.test",
      })
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.isOwnerChanged).toHaveBeenCalledWith(
      true
    )
  })

  it("should respond to SET_MENU_ITEMS message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_MENU_ITEMS",
          items: [{ type: "separator" }],
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.hostMenuItemsChanged
    ).toHaveBeenCalledWith([{ type: "separator" }])
  })

  it("should respond to SET_METADATA message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_METADATA",
          metadata: { hostedAt: "maya", owner: "corgi", repo: "streamlit" },
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.deployedAppMetadataChanged
    ).toHaveBeenCalledWith({
      hostedAt: "maya",
      owner: "corgi",
      repo: "streamlit",
    })
  })

  it("can process a received SET_PAGE_LINK_BASE_URL message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_PAGE_LINK_BASE_URL",
          pageLinkBaseUrl: "https://share.streamlit.io/vdonato/foo/bar",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.pageLinkBaseUrlChanged
    ).toHaveBeenCalledWith("https://share.streamlit.io/vdonato/foo/bar")
  })

  it("can process a received SET_SIDEBAR_CHEVRON_DOWNSHIFT message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_SIDEBAR_CHEVRON_DOWNSHIFT",
          sidebarChevronDownshift: 50,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.sidebarChevronDownshiftChanged
    ).toHaveBeenCalledWith(50)
  })

  it("can process a received SET_SIDEBAR_NAV_VISIBILITY message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_SIDEBAR_NAV_VISIBILITY",
          hidden: true,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.hostHideSidebarNavChanged
    ).toHaveBeenCalledWith(true)
  })

  it("can process a received SET_TOOLBAR_ITEMS message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_TOOLBAR_ITEMS",
          items: [
            {
              borderless: true,
              label: "",
              icon: "star.svg",
              key: "favorite",
            },
          ],
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.hostToolbarItemsChanged
    ).toHaveBeenCalledWith([
      {
        borderless: true,
        icon: "star.svg",
        key: "favorite",
        label: "",
      },
    ])
  })

  it("should respond to UPDATE_HASH message", () => {
    expect(window.location.hash).toEqual("")

    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "UPDATE_HASH",
          hash: "#somehash",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(window.location.hash).toEqual("#somehash")
  })

  it("can process a received UPDATE_FROM_QUERY_PARAMS message", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "UPDATE_FROM_QUERY_PARAMS",
          queryParams: "foo=bar",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.queryParamsChanged).toHaveBeenCalledWith(
      "foo=bar"
    )
    // @ts-expect-error - props are private
    expect(hostCommunicationMgr.props.sendRerunBackMsg).toHaveBeenCalled()
  })

  it("can process a received SET_CUSTOM_THEME_CONFIG message", () => {
    const mockCustomThemeConfig = {
      primaryColor: "#1A6CE7",
      backgroundColor: "#FFFFFF",
      secondaryBackgroundColor: "#F5F5F5",
      textColor: "#1A1D21",
    }
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_CUSTOM_THEME_CONFIG",
          themeInfo: mockCustomThemeConfig,
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.themeChanged
    ).toHaveBeenCalledWith(undefined, mockCustomThemeConfig)
  })

  it("can process a received SET_CUSTOM_THEME_CONFIG message with a dark theme name", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_CUSTOM_THEME_CONFIG",
          themeName: "Dark",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.themeChanged
    ).toHaveBeenCalledWith("Dark", undefined)
  })

  it("can process a received SET_CUSTOM_THEME_CONFIG message with a light theme name", () => {
    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "SET_CUSTOM_THEME_CONFIG",
          themeName: "Light",
        },
        origin: "https://devel.streamlit.test",
      })
    )

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.themeChanged
    ).toHaveBeenCalledWith("Light", undefined)
  })

  it("can process a received SET_FILE_UPLOAD_CLIENT_CONFIG message", () => {
    const message = new MessageEvent("message", {
      data: {
        stCommVersion: HOST_COMM_VERSION,
        type: "SET_FILE_UPLOAD_CLIENT_CONFIG",
        prefix: "https://someprefix.com/hello/",
        headers: {
          header1: "header1value",
          header2: "header2value",
        },
      },
      origin: "https://devel.streamlit.test",
    })
    dispatchEvent("message", message)

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.fileUploadClientConfigChanged
    ).toHaveBeenCalledWith({
      prefix: "https://someprefix.com/hello/",
      headers: {
        header1: "header1value",
        header2: "header2value",
      },
    })
  })

  it("can process a received RESTART_WEBSOCKET_CONNECTION message", () => {
    const message = new MessageEvent("message", {
      data: {
        stCommVersion: HOST_COMM_VERSION,
        type: "RESTART_WEBSOCKET_CONNECTION",
      },
      origin: "https://devel.streamlit.test",
    })
    dispatchEvent("message", message)

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.restartWebsocketConnection
    ).toHaveBeenCalled()
  })

  it("can process a received TERMINATE_WEBSOCKET_CONNECTION message", () => {
    const message = new MessageEvent("message", {
      data: {
        stCommVersion: HOST_COMM_VERSION,
        type: "TERMINATE_WEBSOCKET_CONNECTION",
      },
      origin: "https://devel.streamlit.test",
    })
    dispatchEvent("message", message)

    expect(
      // @ts-expect-error - props are private
      hostCommunicationMgr.props.terminateWebsocketConnection
    ).toHaveBeenCalled()
  })
})

describe("Test different origins", () => {
  let hostCommunicationMgr: HostCommunicationManager
  // eslint-disable-next-line @typescript-eslint/no-explicit-any -- TODO: Replace 'any' with a more specific type.
  let dispatchEvent: any

  beforeEach(() => {
    hostCommunicationMgr = new HostCommunicationManager({
      streamlitExecutionStartedAt: 100,
      themeChanged: vi.fn(),
      sendRerunBackMsg: vi.fn(),
      pageChanged: vi.fn(),
      closeModal: vi.fn(),
      stopScript: vi.fn(),
      rerunScript: vi.fn(),
      clearCache: vi.fn(),
      sendAppHeartbeat: vi.fn(),
      setInputsDisabled: vi.fn(),
      fileUploadClientConfigChanged: vi.fn(),
      isOwnerChanged: vi.fn(),
      hostMenuItemsChanged: vi.fn(),
      hostToolbarItemsChanged: vi.fn(),
      hostHideSidebarNavChanged: vi.fn(),
      sidebarChevronDownshiftChanged: vi.fn(),
      pageLinkBaseUrlChanged: vi.fn(),
      queryParamsChanged: vi.fn(),
      deployedAppMetadataChanged: vi.fn(),
      restartWebsocketConnection: vi.fn(),
      terminateWebsocketConnection: vi.fn(),
    })

    dispatchEvent = mockEventListeners()
  })

  afterEach(() => {
    window.location.hash = ""
  })

  it("exact pattern", () => {
    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://share.streamlit.io"],
      useExternalAuthToken: false,
    })

    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "UPDATE_HASH",
          hash: "#somehash",
        },
        origin: "http://share.streamlit.io",
      })
    )

    expect(window.location.hash).toEqual("#somehash")
  })

  it("wildcard pattern", () => {
    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://*.streamlitapp.com"],
      useExternalAuthToken: false,
    })

    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "UPDATE_HASH",
          hash: "#otherhash",
        },
        origin: "http://cool-cucumber-fa9ds9f.streamlitapp.com",
      })
    )

    expect(window.location.hash).toEqual("#otherhash")
  })

  it("ignores non-matching origins", () => {
    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://share.streamlit.io"],
      useExternalAuthToken: false,
    })

    dispatchEvent(
      "message",
      new MessageEvent("message", {
        data: {
          stCommVersion: HOST_COMM_VERSION,
          type: "UPDATE_HASH",
          hash: "#corgi",
        },
        origin: "http://example.com",
      })
    )

    expect(window.location.hash).toEqual("")
  })
})

describe("HostCommunicationManager external auth token handling", () => {
  let hostCommunicationMgr: HostCommunicationManager

  beforeEach(() => {
    hostCommunicationMgr = new HostCommunicationManager({
      streamlitExecutionStartedAt: 100,
      themeChanged: vi.fn(),
      sendRerunBackMsg: vi.fn(),
      pageChanged: vi.fn(),
      closeModal: vi.fn(),
      stopScript: vi.fn(),
      rerunScript: vi.fn(),
      clearCache: vi.fn(),
      sendAppHeartbeat: vi.fn(),
      setInputsDisabled: vi.fn(),
      fileUploadClientConfigChanged: vi.fn(),
      isOwnerChanged: vi.fn(),
      hostMenuItemsChanged: vi.fn(),
      hostToolbarItemsChanged: vi.fn(),
      hostHideSidebarNavChanged: vi.fn(),
      sidebarChevronDownshiftChanged: vi.fn(),
      pageLinkBaseUrlChanged: vi.fn(),
      queryParamsChanged: vi.fn(),
      deployedAppMetadataChanged: vi.fn(),
      restartWebsocketConnection: vi.fn(),
      terminateWebsocketConnection: vi.fn(),
    })
  })

  it("resolves promise to undefined immediately if useExternalAuthToken is false", async () => {
    const setAllowedOriginsFunc = vi.spyOn(
      hostCommunicationMgr,
      "setAllowedOrigins"
    )

    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://devel.streamlit.test"],
      useExternalAuthToken: false,
    })

    expect(setAllowedOriginsFunc).toHaveBeenCalled()
    // @ts-expect-error - deferredAuthToken is private
    await expect(hostCommunicationMgr.deferredAuthToken.promise).resolves.toBe(
      undefined
    )
  })

  it("waits to receive SET_AUTH_TOKEN message before resolving promise if useExternalAuthToken is true", async () => {
    const dispatchEvent = mockEventListeners()

    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://devel.streamlit.test"],
      useExternalAuthToken: true,
    })
    // Asynchronously send a SET_AUTH_TOKEN message to the
    // HostCommunicationManager, which won't proceed past the `await`
    // statement below until the message is received and handled.
    setTimeout(() => {
      dispatchEvent(
        "message",
        new MessageEvent("message", {
          data: {
            stCommVersion: HOST_COMM_VERSION,
            type: "SET_AUTH_TOKEN",
            authToken: "i am an auth token",
          },
          origin: "http://devel.streamlit.test",
        })
      )
    })

    // @ts-expect-error - deferredAuthToken is private
    await expect(hostCommunicationMgr.deferredAuthToken.promise).resolves.toBe(
      "i am an auth token"
    )

    // Reset the auth token and do everything again to confirm that we don't
    // incorrectly resolve to an old value after resetAuthToken is called.
    hostCommunicationMgr.resetAuthToken()

    // Simulate the browser tab disconnecting and reconnecting, which from the
    // HostCommunication's perspective is only seen as a new call to
    // setAllowedOrigins.
    hostCommunicationMgr.setAllowedOrigins({
      allowedOrigins: ["http://devel.streamlit.test"],
      useExternalAuthToken: true,
    })

    setTimeout(() => {
      dispatchEvent(
        "message",
        new MessageEvent("message", {
          data: {
            stCommVersion: HOST_COMM_VERSION,
            type: "SET_AUTH_TOKEN",
            authToken: "i am a NEW auth token",
          },
          origin: "http://devel.streamlit.test",
        })
      )
    })

    // @ts-expect-error - deferredAuthToken is private
    await expect(hostCommunicationMgr.deferredAuthToken.promise).resolves.toBe(
      "i am a NEW auth token"
    )
  })
})