baseTest.ts

import type { ElectronApplication, Page, TestInfo } from "@playwright/test";
import { _electron as electron, expect, test as testBase } from "@playwright/test";
import archiver from "archiver";
import { stubAllDialogs, stubDialog } from "electron-playwright-helpers";
import { createWriteStream, existsSync, mkdtempSync, readFileSync } from "fs";
import { tmpdir } from "os";
import path from "path";
import { DEBUG_LOGGING_ENABLED } from "./constants";
import { NotificationArea } from "./objects/notifications/NotificationArea";
import { Quickpick } from "./objects/quickInputs/Quickpick";
import {
  DEFAULT_CCLOUD_TOPIC_REPLICATION_FACTOR,
  SelectKafkaCluster,
  TopicsView,
} from "./objects/views/TopicsView";
import type { CCloudConnectionItem } from "./objects/views/viewItems/CCloudConnectionItem";
import type { DirectConnectionItem } from "./objects/views/viewItems/DirectConnectionItem";
import type { LocalConnectionItem } from "./objects/views/viewItems/LocalConnectionItem";
import type { DirectConnectionOptions, LocalConnectionOptions } from "./types/connection";
import { ConnectionType, FormConnectionType, SupportedAuthType } from "./types/connection";
import type { TopicConfig } from "./types/topic";
import { executeVSCodeCommand } from "./utils/commands";
import {
  setupCCloudConnection,
  setupDirectConnection,
  setupLocalConnection,
  teardownLocalConnection,
} from "./utils/connections";
import { produceMessages } from "./utils/producer";
import { configureVSCodeSettings } from "./utils/settings";
import { randomHexString } from "./utils/strings";
import { openConfluentSidebar, prepareTestWorkspace } from "./utils/workspace";

// NOTE: we can't import these two directly from 'global.setup.ts'
// cached test setup file path that's shared across worker processes
const TEST_SETUP_CACHE_FILE = path.join(tmpdir(), "vscode-e2e-test-setup-cache.json");

interface TestSetupCache {
  vscodeExecutablePath: string;
  outPath: string;
}

/** Get the test setup cache created by the global setup, avoiding repeated VS Code setup logging. */
function getTestSetupCache(): TestSetupCache {
  if (!existsSync(TEST_SETUP_CACHE_FILE)) {
    throw new Error(`Test setup cache file not found at ${TEST_SETUP_CACHE_FILE}.`);
  }
  try {
    const cacheContent = readFileSync(TEST_SETUP_CACHE_FILE, "utf-8");
    return JSON.parse(cacheContent);
  } catch (error) {
    throw new Error(`Failed to read test setup cache: ${error}`);
  }
}

interface VSCodeFixtures {
  /** Temporary directory for the test to use. */
  testTempDir: string;

  /** The launched Electron application (VS Code). */
  electronApp: ElectronApplication;
  /** The first window of the launched Electron application (VS Code). */
  page: Page;

  /** Auto fixture to add annotations needed for test results analysis. */
  annotatePlatform: void;

  /** Open the Confluent view container from the primary sidebar, activating the extension if necessary. */
  openExtensionSidebar: void;

  /**
   * Connection type to set up for parameterized tests.
   * Used by the `connectionItem` fixture to determine which connection to set up.
   */
  connectionType: ConnectionType | undefined;
  /**
   * Configuration options for setting up a direct connection with the {@linkcode directConnection} fixture.
   */
  directConnectionConfig: DirectConnectionOptions;
  /**
   * Configuration options for setting up a local connection with the {@linkcode localConnection} fixture.
   * If not provided, the local connection will set up both Kafka and Schema Registry by default.
   */
  localConnectionConfig: LocalConnectionOptions;
  /**
   * Set up a connection based on the {@linkcode connectionType} option and returns the associated
   * connection item ({@link CCloudConnectionItem}, {@link DirectConnectionItem}, or {@link LocalConnectionItem}).
   */
  connectionItem: CCloudConnectionItem | DirectConnectionItem | LocalConnectionItem;

  /**
   * Configuration options for creating a topic with the {@linkcode topic} fixture.
   */
  topicConfig: TopicConfig | undefined;
  /**
   * Set up a topic based on the {@linkcode topicConfig} option and return the associated topic
   * `name` for tests to reference.
   */
  topic: string;
}

export const test = testBase.extend<VSCodeFixtures>({
  testTempDir: async ({}, use) => {
    const tempDir = mkdtempSync(path.join(tmpdir(), "vscode-test-"));

    await use(tempDir);
  },

  electronApp: async ({ testTempDir }, use, testInfo) => {
    const testConfigs = getTestSetupCache();

    // launch VS Code with Electron using args pattern from vscode-test
    const electronApp = await electron.launch({
      executablePath: testConfigs.vscodeExecutablePath,
      args: [
        // same as the Mocha test args in Gulpfile.js:
        "--no-sandbox",
        "--skip-release-notes",
        "--skip-welcome",
        "--disable-gpu",
        "--disable-updates",
        "--disable-workspace-trust",
        "--disable-extensions",
        // required to prevent test resources being saved to user's real profile
        `--user-data-dir=${testTempDir}`,
        // additional args needed for the Electron launch:
        `--extensionDevelopmentPath=${testConfigs.outPath}`,
      ],
    });
    if (!electronApp) {
      throw new Error("Failed to launch VS Code electron app");
    }

    const context = electronApp.context();
    // grant clipboard access for writing content into editors and reading copied values
    await context.grantPermissions(["clipboard-read", "clipboard-write"]);
    // always start tracing manually, but decide later whether to save it based on test result
    await context.tracing.start({
      screenshots: true,
      snapshots: true,
      sources: true,
      title: `${process.platform} ${process.arch}: ${testInfo.title} (${testInfo.tags.join(", ")})`,
    });

    try {
      // wait for VS Code to be ready before trying to stub dialogs
      const page = await electronApp.firstWindow();
      await page.waitForLoadState("domcontentloaded");
      await page.locator(".monaco-workbench").waitFor({ timeout: 30000 });

      // Stub all dialogs by default; tests can still override as needed.
      // For available `method` values to use with `stubMultipleDialogs`, see:
      // https://www.electronjs.org/docs/latest/api/dialog
      await stubAllDialogs(electronApp);

      await use(electronApp);
    } finally {
      // only save and attach the trace for failed tests
      if (testInfo.status !== testInfo.expectedStatus) {
        const tracePath = path.join(testInfo.outputDir, "trace.zip");
        await context.tracing.stop({ path: tracePath });
        await testInfo.attach("trace", { path: tracePath, contentType: "application/zip" });
      } else {
        await context.tracing.stop();
      }
    }

    await shutdownElectronApp(electronApp);
  },

  page: async ({ electronApp, testTempDir }, use, testInfo) => {
    const page = await electronApp.firstWindow();

    await globalBeforeEach(page, testTempDir);

    await use(page);

    await globalAfterEach(testTempDir, electronApp, page, testInfo);
  },

  annotatePlatform: [
    async ({}, use, testInfo) => {
      testInfo.annotations.push({
        type: "platform",
        description: `${process.platform} ${process.arch}`,
      });
      await use();
    },
    { auto: true },
  ],

  openExtensionSidebar: [
    async ({ page }, use) => {
      await openConfluentSidebar(page);

      await use();

      // no explicit teardown needed
    },
    { auto: true }, // automatically run for all tests unless opted out
  ],

  directConnectionConfig: [
    {
      formConnectionType: FormConnectionType.ConfluentCloud,
      kafkaConfig: {
        bootstrapServers: process.env.E2E_KAFKA_BOOTSTRAP_SERVERS!,
        authType: SupportedAuthType.API,
        credentials: {
          api_key: process.env.E2E_KAFKA_API_KEY!,
          api_secret: process.env.E2E_KAFKA_API_SECRET!,
        },
      },
      schemaRegistryConfig: {
        uri: process.env.E2E_SR_URL!,
        authType: SupportedAuthType.API,
        credentials: {
          api_key: process.env.E2E_SR_API_KEY!,
          api_secret: process.env.E2E_SR_API_SECRET!,
        },
      },
    },
    { option: true },
  ],

  localConnectionConfig: [{ schemaRegistry: true }, { option: true }],

  // no default value, must be provided by test
  connectionType: [undefined, { option: true }],

  connectionItem: async (
    { electronApp, page, connectionType, directConnectionConfig, localConnectionConfig },
    use,
    testInfo,
  ) => {
    if (!connectionType) {
      throw new Error(
        "connectionType must be set, like `test.use({ connectionType: ConnectionType.Ccloud })`",
      );
    }

    let connection: CCloudConnectionItem | DirectConnectionItem | LocalConnectionItem;

    // setup
    switch (connectionType) {
      case ConnectionType.Ccloud:
        connection = await setupCCloudConnection(
          page,
          electronApp,
          process.env.E2E_USERNAME!,
          process.env.E2E_PASSWORD!,
          testInfo,
        );
        break;
      case ConnectionType.Direct:
        connection = await setupDirectConnection(page, directConnectionConfig);
        break;
      case ConnectionType.Local:
        connection = await setupLocalConnection(page, localConnectionConfig);
        break;
      default:
        throw new Error(`Unsupported connection type: ${connectionType}`);
    }

    await use(connection);

    // teardown
    switch (connectionType) {
      case ConnectionType.Ccloud:
        // no explicit teardown needed since shutting down the extension+sidecar will invalidate the
        // CCloud auth session
        break;
      case ConnectionType.Direct:
        // no teardown needed since each test will use its own storage in TMPDIR, so any direct
        // connections created will be cleaned up automatically, and subsequent tests will use their
        // own blank-slate storage
        break;
      case ConnectionType.Local:
        // local resources are discovered automatically through the Docker engine API, so we need
        // to explicitly stop them to ensure the next tests can start them fresh
        await teardownLocalConnection(connection.page, {
          schemaRegistry: localConnectionConfig.schemaRegistry,
        });
        break;
      default:
        throw new Error(`Unsupported connection type: ${connectionType}`);
    }
  },

  // no default value, must be provided by test
  topicConfig: [undefined, { option: true }],

  topic: async (
    { page, connectionType, connectionItem, topicConfig, directConnectionConfig },
    use,
  ) => {
    if (!connectionType) {
      throw new Error(
        "connectionType must be set, like `test.use({ connectionType: ConnectionType.Ccloud })`",
      );
    }
    if (!topicConfig) {
      throw new Error(
        "topicConfig must be set, like `test.use({ topicConfig: { name: 'my-topic' } })`",
      );
    }

    // ensure connection has resources available to work with
    await expect(connectionItem.locator).toHaveAttribute("aria-expanded", "true");

    const topicName = `${topicConfig.name}-${randomHexString(6)}`;

    // set default replication factor (if it wasn't provided) based on connection type
    const replicationFactor =
      topicConfig.replicationFactor ??
      (connectionType === ConnectionType.Local ? 1 : DEFAULT_CCLOUD_TOPIC_REPLICATION_FACTOR);

    const numPartitions = topicConfig.numPartitions ?? 1;

    // if we need to produce messages, we likely have an API key/secret we need to match to a
    // specific CCloud cluster, so we can't use the first one that shows up in the resources view
    // (LOCAL/DIRECT connections don't have multiple clusters, so we can just skip this)
    let clusterLabel: string | RegExp | undefined;
    if (topicConfig.produce && connectionType === ConnectionType.Ccloud) {
      clusterLabel = topicConfig.clusterLabel ?? process.env.E2E_KAFKA_CLUSTER_NAME!;
    }

    // setup: create the topic
    const topicsView = new TopicsView(page);
    await topicsView.loadTopics(connectionType, SelectKafkaCluster.FromResourcesView, clusterLabel);
    await topicsView.createTopic(topicName, numPartitions, replicationFactor);
    await topicsView.getTopicItem(topicName); // verify topic shows up in the view

    // produce messages to the topic if specified
    if (topicConfig.produce) {
      await produceMessages(
        page,
        connectionType,
        topicName,
        topicConfig.produce,
        directConnectionConfig,
      );
    }

    await use(topicName);

    // teardown: delete the topic
    // (explicitly make sure the sidebar is open and we reload the topics view in the event a test
    // navigated away to a new window or sidebar)
    await openConfluentSidebar(page);
    await topicsView.loadTopics(connectionType, SelectKafkaCluster.FromResourcesView, clusterLabel);
    await topicsView.deleteTopic(topicName);
  },
});

/**
 * Global setup that runs before each test.
 *
 * NOTE: Due to our Electron launch setup, this is more reliable than using
 * {@linkcode https://playwright.dev/docs/api/class-test#test-before-each test.beforeEach()}, which
 * did not consistently run before each test.
 */
async function globalBeforeEach(page: Page, testTempDir: string): Promise<void> {
  // make sure settings are set to defaults for each test
  configureVSCodeSettings(testTempDir);

  // dismiss the disabled-extensions toast and collapse the secondary sidebar
  await prepareTestWorkspace(page);
}

async function globalAfterEach(
  testTempDir: string,
  electronApp: ElectronApplication,
  page: Page,
  testInfo: TestInfo,
): Promise<void> {
  // try to save extension and sidecar logs and attach them to the results for each test, but don't
  // fail tests (that would otherwise pass) if either time out since they're nice-to-haves
  await saveExtensionLogs(testTempDir, electronApp, page, testInfo);
  await saveSidecarLogs(testTempDir, electronApp, page, testInfo);
  // also include any additional logs from the VS Code window itself (main, window, extension host, etc)
  await saveVSCodeWindowLogs(testTempDir, testInfo);
}

async function saveExtensionLogs(
  testTempDir: string,
  electronApp: ElectronApplication,
  page: Page,
  testInfo: TestInfo,
): Promise<void> {
  const extensionLogPath = path.join(testTempDir, "vscode-confluent.log");
  await stubDialog(electronApp, "showSaveDialog", {
    filePath: extensionLogPath,
  });
  await executeVSCodeCommand(page, "confluent.support.saveLogs");

  try {
    // wait for info notification indicating extension log file was saved
    const notificationArea = new NotificationArea(page);
    const extLogSuccess = notificationArea.infoNotifications.filter({
      hasText: "Confluent extension log file saved successfully.",
    });
    await expect(extLogSuccess).toHaveCount(1, { timeout: 5000 });
    // attach the extension log to the test results
    await testInfo.attach("vscode-confluent.log", {
      path: extensionLogPath,
      contentType: "text/plain",
    });
  } catch {
    console.error("Failed to save extension logs");
  }
}

async function saveSidecarLogs(
  testTempDir: string,
  electronApp: ElectronApplication,
  page: Page,
  testInfo: TestInfo,
): Promise<void> {
  const sidecarLogPath = path.join(testTempDir, "vscode-confluent-sidecar.log");
  await stubDialog(electronApp, "showSaveDialog", {
    filePath: sidecarLogPath,
  });
  await executeVSCodeCommand(page, "confluent.support.saveSidecarLogs");

  // select the formatted log option in the quick pick
  const formatQuickPick = new Quickpick(page);
  await expect(formatQuickPick.locator).toBeVisible({ timeout: 5000 });
  await formatQuickPick.selectItemByText("Human-readable format");

  try {
    // wait for info notification indicating sidecar log file was saved
    const notificationArea = new NotificationArea(page);
    const sidecarLogSuccess = notificationArea.infoNotifications.filter({
      hasText: "Confluent extension sidecar log file saved successfully.",
    });
    await expect(sidecarLogSuccess).toHaveCount(1, { timeout: 5000 });
    // attach the sidecar log to the test results
    await testInfo.attach("vscode-confluent-sidecar.log", {
      path: sidecarLogPath,
      contentType: "text/plain",
    });
  } catch {
    console.error("Failed to save sidecar logs");
  }
}

async function saveVSCodeWindowLogs(testTempDir: string, testInfo: TestInfo): Promise<void> {
  // this will end up looking like `${testTempDir}/vscode-test-<hash>/logs/YYYYmmddTHHMMSS/window1/`
  // which then has `exthost/` and `output_YYYYmmddTHHMMSS/` subdirectories
  const logsDir = path.join(testTempDir, "logs");
  if (!existsSync(logsDir)) {
    console.warn("VS Code logs directory does not exist:", logsDir);
    return;
  }

  const zipFileName = "vscode-window-logs.zip";
  try {
    const zipPath = path.join(testTempDir, zipFileName);
    const output = createWriteStream(zipPath);
    const archive = archiver("zip", { zlib: { level: 9 } });

    // wait for the archive to finish
    await new Promise<void>((resolve, reject) => {
      output.on("close", () => resolve());
      archive.on("error", (err) => reject(err));
      archive.pipe(output);
      archive.directory(logsDir, false);
      archive.finalize();
    });

    await testInfo.attach(zipFileName, {
      path: zipPath,
      contentType: "application/zip",
    });
  } catch (error) {
    console.error("Error zipping VS Code logs directory:", error);
  }
}

/** Partial type for handles that may have an unref() method. */
type PartialHandle = { unref?: () => void };

/**
 * Extended type for process with the (deprecated) `_getActiveHandles` method.
 *
 * NOTE: https://nodejs.org/api/deprecations.html#DEP0161 suggests using
 * `process.getActiveResourcesInfo()`, but that only provides the names of active resources,
 * not the actual handles to unref and allow the process to exit cleanly.
 */
type NodeProcessWithGetActiveHandles = NodeJS.Process & {
  _getActiveHandles?: () => PartialHandle[];
};

/**
 * Shut down the Electron app and clean up any lingering handles/requests so the Playwright worker
 * can exit cleanly.
 *
 * With Playwright v1.50.0+, the behavior of `electronApp.close()` changes so it waits for the
 * Electron process to exit, but if there are any lingering handles or requests, that may never
 * happen, causing the test worker to hang and eventually time out.
 *
 * This function attempts to close the app gracefully, but if it doesn't exit within a timeout,
 * it force kills the process group. It also unrefs any remaining active handles to allow the
 * worker teardown to complete.
 */
async function shutdownElectronApp(electronApp: ElectronApplication): Promise<void> {
  // 1) wait for close() to complete, but with a short timeout
  let timeout: NodeJS.Timeout | undefined;
  try {
    await Promise.race([
      electronApp.close(),
      new Promise((resolve) => {
        timeout = setTimeout(resolve, 1000);
      }),
    ]);
  } catch {
    // no need to deal with errors when we're trying to shut everything down here
  } finally {
    if (timeout !== undefined) {
      clearTimeout(timeout);
    }
  }

  // 2) check if the Electron process is still running, and force kill if so
  try {
    const proc = electronApp.process();
    const pid = proc?.pid;
    if (pid && pid > 1) {
      process.kill(pid, 0); // status check
      console.warn("Electron still running after close(), force killing process group...");
      process.kill(-pid, 9); // SIGKILL entire process group (negative PID)
    }
  } catch {
    // process no longer exists, no action needed
  }

  // 3) unref any remaining handles to allow the worker teardown to complete cleanly
  try {
    const handles = (process as NodeProcessWithGetActiveHandles)._getActiveHandles?.() || [];
    let unrefdHandles = 0;
    handles.forEach((handle: PartialHandle) => {
      if (handle && typeof handle.unref === "function") {
        try {
          // see description for https://nodejs.org/api/process.html#processunrefmayberefable
          handle.unref();
          unrefdHandles++;
        } catch {
          // some handles may not support unref, not much we can do
        }
      }
    });
    if (DEBUG_LOGGING_ENABLED && handles.length) {
      console.debug(`Unreferenced ${unrefdHandles}/${handles.length} handle(s)`);
    }
  } catch {
    // ignore errors during unref since we can't do much about them
  }
}