cosmos-explorer/src/HostedExplorer.test.tsx

jest.mock("./hooks/useAADAuth");
jest.mock("./hooks/useConfig");
jest.mock("./hooks/usePortalAccessToken");
jest.mock("./Platform/Hosted/Components/ConnectExplorer");
jest.mock("./Shared/appInsights");
jest.mock("./Platform/Hosted/Components/AccountSwitcher", () => ({
  AccountSwitcher: () => null,
}));
jest.mock("./Platform/Hosted/Components/DirectoryPickerPanel", () => ({
  DirectoryPickerPanel: () => null,
}));
jest.mock("./Platform/Hosted/Components/FeedbackCommandButton", () => ({
  FeedbackCommandButton: () => null,
}));
jest.mock("./Platform/Hosted/Components/MeControl", () => ({
  MeControl: () => null,
}));
jest.mock("./Platform/Hosted/Components/SignInButton", () => ({
  SignInButton: () => null,
}));
jest.mock("./Platform/Hosted/Components/AadAuthorizationFailure", () => ({
  AadAuthorizationFailure: () => null,
}));

import "@testing-library/jest-dom";
import { act, render } from "@testing-library/react";
import React from "react";
import { useAADAuth } from "./hooks/useAADAuth";
import { useConfig } from "./hooks/useConfig";
import { useTokenMetadata } from "./hooks/usePortalAccessToken";
import { App } from "./HostedExplorer";
import {
  ConnectExplorer,
  fetchEncryptedToken,
  validateDirectConnectionStringConnectivity,
} from "./Platform/Hosted/Components/ConnectExplorer";

const mockFetchEncryptedToken = fetchEncryptedToken as jest.MockedFunction<typeof fetchEncryptedToken>;
const mockValidateDirectConnectionStringConnectivity =
  validateDirectConnectionStringConnectivity as jest.MockedFunction<typeof validateDirectConnectionStringConnectivity>;

(ConnectExplorer as jest.Mock).mockImplementation(() => <div data-testid="connect-explorer" />);

import type { AccountInfo } from "@azure/msal-browser";
import type { AadAuthFailure } from "./hooks/useAADAuth";

const defaultAADAuth = {
  isLoggedIn: false,
  armToken: "",
  graphToken: "",
  account: undefined as AccountInfo | null | undefined,
  tenantId: "",
  logout: jest.fn(),
  login: jest.fn(),
  switchTenant: jest.fn(),
  authFailure: undefined as AadAuthFailure | null | undefined,
};

beforeEach(() => {
  jest.clearAllMocks();
  (useAADAuth as jest.Mock).mockReturnValue(defaultAADAuth);
  (useConfig as jest.Mock).mockReturnValue({});
  (useTokenMetadata as jest.Mock).mockReturnValue(undefined);
  mockFetchEncryptedToken.mockResolvedValue("encrypted-token");
  mockValidateDirectConnectionStringConnectivity.mockResolvedValue(undefined);
});

const dispatchPostMessage = (data: unknown, origin: string) => {
  const event = new MessageEvent("message", { data, origin });
  window.dispatchEvent(event);
};

// Deliberately invalid account name
const FAKE_ACCOUNT_NAME: string = "-FakeAccount-";
const FAKE_KEY: string = "<redacted-test-key>";

describe("HostedExplorer tryCosmosDB postMessage handler", () => {
  it("signs a valid SQL connection string client-side without calling the backend", async () => {
    render(<App />);

    const validConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.documents.azure.com:443/;AccountKey=${FAKE_KEY};`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: validConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
    expect(mockValidateDirectConnectionStringConnectivity).toHaveBeenCalled();
  });

  it("accepts a valid Mongo connection string from an allowed origin", async () => {
    render(<App />);

    const mongoConnStr = `mongodb://${FAKE_ACCOUNT_NAME}:${FAKE_KEY}@${FAKE_ACCOUNT_NAME}.documents.azure.com:10255`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: mongoConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockFetchEncryptedToken).toHaveBeenCalledWith(mongoConnStr);
  });

  it("accepts a valid Cassandra connection string from an allowed origin", async () => {
    render(<App />);

    const cassandraConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.cassandra.cosmosdb.azure.com:443/;AccountKey=${FAKE_KEY};`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: cassandraConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockFetchEncryptedToken).toHaveBeenCalledWith(cassandraConnStr);
  });

  it("signs a valid Table connection string client-side without calling the backend", async () => {
    render(<App />);

    const tableConnStr = `DefaultEndpointsProtocol=https;AccountName=${FAKE_ACCOUNT_NAME};AccountKey=${FAKE_KEY};TableEndpoint=https://${FAKE_ACCOUNT_NAME}.table.cosmosdb.azure.com:443/;`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: tableConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
    expect(mockValidateDirectConnectionStringConnectivity).toHaveBeenCalled();
  });

  it("signs a valid Gremlin connection string client-side without calling the backend", async () => {
    render(<App />);

    const gremlinConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.documents.azure.com:443/;AccountKey=${FAKE_KEY};ApiKind=Gremlin;`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: gremlinConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
    expect(mockValidateDirectConnectionStringConnectivity).toHaveBeenCalled();
  });

  it("does not open the Data Explorer when the Cosmos client cannot connect", async () => {
    mockValidateDirectConnectionStringConnectivity.mockRejectedValue(new Error("Unable to connect to the account."));
    const { container } = render(<App />);

    const validConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.documents.azure.com:443/;AccountKey=${FAKE_KEY};`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: validConnStr },
        "https://cosmos.azure.com",
      );
      await Promise.resolve();
    });

    expect(mockValidateDirectConnectionStringConnectivity).toHaveBeenCalled();
    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
    expect(container.querySelector('[data-test="DataExplorerFrame"]')).toBeNull();
  });

  it("rejects messages from a disallowed origin", async () => {
    render(<App />);

    const validConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.documents.azure.com:443/;AccountKey=${FAKE_KEY};`;

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: validConnStr },
        "https://evil.example.com",
      );
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
  });

  it("rejects messages with an invalid connection string format", async () => {
    render(<App />);

    await act(async () => {
      dispatchPostMessage(
        { type: "tryCosmosDBConnectionString", connectionString: "not-a-real-connection-string" },
        "https://cosmos.azure.com",
      );
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
  });

  it("rejects messages with a non-string connection string value", async () => {
    render(<App />);

    await act(async () => {
      dispatchPostMessage({ type: "tryCosmosDBConnectionString", connectionString: 12345 }, "https://cosmos.azure.com");
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
  });

  it("rejects messages with a missing connection string", async () => {
    render(<App />);

    await act(async () => {
      dispatchPostMessage({ type: "tryCosmosDBConnectionString" }, "https://cosmos.azure.com");
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
  });

  it("ignores messages with an unrelated type", async () => {
    render(<App />);

    const validConnStr = `AccountEndpoint=https://${FAKE_ACCOUNT_NAME}.documents.azure.com:443/;AccountKey=${FAKE_KEY};`;

    await act(async () => {
      dispatchPostMessage({ type: "someOtherMessage", connectionString: validConnStr }, "https://cosmos.azure.com");
    });

    expect(mockFetchEncryptedToken).not.toHaveBeenCalled();
  });

  it("sends tryCosmosDBReady to opener when present", () => {
    const mockPostMessage = jest.fn();
    const originalOpener = window.opener;
    Object.defineProperty(window, "opener", {
      value: { postMessage: mockPostMessage },
      writable: true,
      configurable: true,
    });

    render(<App />);

    expect(mockPostMessage).toHaveBeenCalledWith({ type: "tryCosmosDBReady" }, "https://cosmos.azure.com");

    Object.defineProperty(window, "opener", {
      value: originalOpener,
      writable: true,
      configurable: true,
    });
  });

  it("does not crash when there is no opener", () => {
    const originalOpener = window.opener;
    Object.defineProperty(window, "opener", {
      value: null,
      writable: true,
      configurable: true,
    });

    expect(() => render(<App />)).not.toThrow();

    Object.defineProperty(window, "opener", {
      value: originalOpener,
      writable: true,
      configurable: true,
    });
  });
});