SaveState improvements

src/browser/commands/getState.ts

@@ -0,0 +1,25 @@
+import fs from "fs-extra";
+
+import type { Browser } from "../types";
+import type { SaveStateData } from "./saveState";
+import type { StateOpts } from "../../config/types";
+import * as logger from "../../utils/logger";
+
+export type GetStateOptions = Pick<StateOpts, "path">;
+
+export default (browser: Browser): void => {
+    const { publicAPI: session } = browser;
+
+    session.addCommand(
+        "getState",
+        async (options: GetStateOptions = browser.config.stateOpts || {}): Promise<SaveStateData | null> => {
+            if (options.path) {
+                return fs.readJson(options.path);
+            }
+
+            logger.error("Please provide the path to the state file for getState");
+
+            return null;
+        },
+    );
+};

src/browser/commands/index.ts

@@ -12,5 +12,6 @@ export const customCommandFileNames = [
     "waitForStaticToLoad",
     "saveState",
     "restoreState",
+    "getState",
     "unstable_getCdp",
 ];

src/browser/commands/restoreState/index.ts

@@ -2,36 +2,29 @@ import fs from "fs-extra";
 
 import { restoreStorage } from "./restoreStorage";
 
-import * as logger from "../../../utils/logger";
 import type { Browser } from "../../types";
 import { DEVTOOLS_PROTOCOL, WEBDRIVER_PROTOCOL } from "../../../constants/config";
-import {
-    defaultOptions,
-    getOverridesProtocol,
-    getWebdriverFrames,
-    SaveStateData,
-    SaveStateOptions,
-} from "../saveState";
+import { getOverridesProtocol, getWebdriverFrames, SaveStateData } from "../saveState";
 import { getActivePuppeteerPage } from "../../existing-browser";
 import { Cookie } from "@testplane/wdio-protocols";
+import { StateOpts } from "../../../config/types";
 
-export type RestoreStateOptions = SaveStateOptions & {
+export type RestoreStateOptions = Omit<StateOpts, "keepFile"> & {
     data?: SaveStateData;
     refresh?: boolean;
 };
 
 export default (browser: Browser): void => {
     const { publicAPI: session } = browser;
 
-    session.addCommand("restoreState", async (_options: RestoreStateOptions) => {
+    session.addCommand("restoreState", async (_options: RestoreStateOptions = {}) => {
         const currentUrl = new URL(await session.getUrl());
 
         if (!currentUrl.origin || currentUrl.origin === "null") {
-            logger.error("Before restoreState first open page using url command");
-            process.exit(1);
+            throw new Error("Before restoreState first open page using url command");
         }
 
-        const options = { ...defaultOptions, refresh: true, ..._options };
+        const options = { ...browser.config.stateOpts, refresh: true, ..._options };
 
         let restoreState: SaveStateData | undefined = options.data;
 
@@ -40,8 +33,7 @@ export default (browser: Browser): void => {
         }
 
         if (!restoreState) {
-            logger.error("Can't restore state: please provide a path to file or data");
-            process.exit(1);
+            throw new Error("Can't restore state: please provide a path to file or data");
         }
 
         if (restoreState?.cookies && options.cookieFilter) {

src/browser/commands/saveState/index.ts

@@ -7,16 +7,9 @@ import { ExistingBrowser, getActivePuppeteerPage } from "../../existing-browser"
 import * as logger from "../../../utils/logger";
 import { Cookie } from "../../../types";
 import type { Browser } from "../../types";
-
-export type SaveStateOptions = {
-    path?: string;
-
-    cookies?: boolean;
-    localStorage?: boolean;
-    sessionStorage?: boolean;
-
-    cookieFilter?: (cookie: Cookie) => boolean;
-};
+import { MasterEvents } from "../../../events";
+import { StateOpts } from "../../../config/types";
+import { addGlobalFileToRemove } from "../../../globalFilesToRemove";
 
 export type FrameData = StorageData;
 
@@ -25,12 +18,6 @@ export type SaveStateData = {
     framesData: Record<string, FrameData>;
 };
 
-export const defaultOptions = {
-    cookies: true,
-    localStorage: true,
-    sessionStorage: true,
-};
-
 // in case when we use webdriver protocol, bidi and isolation
 // we have to force change protocol to devtools, for use puppeteer,
 // because we use it for create incognito window
@@ -49,15 +36,14 @@ export const getWebdriverFrames = async (session: WebdriverIO.Browser): Promise<
 export default (browser: ExistingBrowser): void => {
     const { publicAPI: session } = browser;
 
-    session.addCommand("saveState", async (_options: SaveStateOptions = {}): Promise<SaveStateData | undefined> => {
+    session.addCommand("saveState", async (_options: StateOpts = {}): Promise<SaveStateData | undefined> => {
         const currentUrl = new URL(await session.getUrl());
 
         if (!currentUrl.origin || currentUrl.origin === "null") {
-            logger.error("Before saveState first open page using url command");
-            process.exit(1);
+            throw new Error("Before saveState first open page using url command");
         }
 
-        const options = { ...defaultOptions, ..._options };
+        const options = { ...browser.config.stateOpts, ..._options };
 
         const data: SaveStateData = {
             framesData: {},
@@ -178,8 +164,27 @@ export default (browser: ExistingBrowser): void => {
             data.cookies = data.cookies.filter(options.cookieFilter);
         }
 
-        if (options && options.path) {
-            await fs.writeJson(options.path, data, { spaces: 2 });
+        const dataIsEmpty = data.cookies?.length === 0 && _.isEmpty(data.framesData);
+
+        if (options && options.path && !dataIsEmpty) {
+            await fs.outputJson(options.path, data, { spaces: 2 });
+
+            if (options.keepFile) {
+                logger.warn(
+                    "\x1b[31mOption keepFile in stateOpts now is true. Please be aware that the file containing authorization data will not be automatically deleted after the tests are completed!\x1b[0m",
+                );
+            } else {
+                if (process.send) {
+                    process.send({
+                        event: MasterEvents.ADD_FILE_TO_REMOVE,
+                        data: options.path,
+                    });
+                }
+
+                addGlobalFileToRemove(options.path);
+
+                browser.emitter.emit(MasterEvents.ADD_FILE_TO_REMOVE, options.path);
+            }
         }
 
         return data;

src/browser/standalone/attachToBrowser.ts

@@ -1,9 +1,11 @@
 import { Config } from "../../config";
 import { ExistingBrowser } from "./../existing-browser";
 import { Calibrator } from "./../calibrator";
-import { AsyncEmitter } from "../../events";
+import { AsyncEmitter, MasterEvents } from "../../events";
 import { BrowserName, type W3CBrowserName, type SessionOptions } from "./../types";
 import { getNormalizedBrowserName } from "../../utils/browser";
+import fs from "fs-extra";
+import { hasGlobalFilesToRemove } from "../../globalFilesToRemove";
 
 export async function attachToBrowser(session: SessionOptions): Promise<WebdriverIO.Browser> {
     const browserName = session.sessionCaps?.browserName || BrowserName.CHROME;
@@ -24,6 +26,8 @@ export async function attachToBrowser(session: SessionOptions): Promise<Webdrive
         },
     };
 
+    const filesToRemove: string[] = [];
+
     const config = new Config({
         browsers: {
             [browserName]: browserConfig,
@@ -36,6 +40,10 @@ export async function attachToBrowser(session: SessionOptions): Promise<Webdrive
 
     const emitter = new AsyncEmitter();
 
+    emitter.on(MasterEvents.ADD_FILE_TO_REMOVE, (path: string) => {
+        filesToRemove.push(path);
+    });
+
     const existingBrowser = new ExistingBrowser(config, {
         id: browserName,
         version: session.sessionCaps?.browserVersion,
@@ -54,6 +62,10 @@ export async function attachToBrowser(session: SessionOptions): Promise<Webdrive
         if (session.driverPid) {
             process.kill(session.driverPid, 9);
         }
+
+        if (filesToRemove.length > 0 && !hasGlobalFilesToRemove()) {
+            await Promise.all(filesToRemove.map(path => fs.remove(path)));
+        }
     });
 
     return existingBrowser.publicAPI;

src/browser/standalone/launchBrowser.ts

@@ -2,12 +2,14 @@ import { Config } from "../../config";
 import { NewBrowser } from "./../new-browser";
 import { ExistingBrowser } from "./../existing-browser";
 import { Calibrator } from "./../calibrator";
-import { AsyncEmitter } from "../../events";
+import { AsyncEmitter, MasterEvents } from "../../events";
 import { BrowserName, type W3CBrowserName } from "./../types";
 import { getNormalizedBrowserName } from "../../utils/browser";
 import { LOCAL_GRID_URL } from "../../constants/config";
 import { WebdriverPool } from "../../browser-pool/webdriver-pool";
 import type { StandaloneBrowserOptionsInput } from "./types";
+import fs from "fs-extra";
+import { hasGlobalFilesToRemove } from "../../globalFilesToRemove";
 
 const webdriverPool = new WebdriverPool();
 
@@ -48,8 +50,11 @@ export async function launchBrowser(
         user: options.user,
         key: options.key,
         prepareBrowser: options.prepareBrowser,
+        stateOpts: options.stateOpts,
     };
 
+    const filesToRemove: string[] = [];
+
     const config = new Config({
         browsers: {
             [browserName]: browserConfig,
@@ -62,6 +67,10 @@ export async function launchBrowser(
 
     const emitter = new AsyncEmitter();
 
+    emitter.on(MasterEvents.ADD_FILE_TO_REMOVE, (path: string) => {
+        filesToRemove.push(path);
+    });
+
     const newBrowser = new NewBrowser(config, {
         id: browserName,
         version: desiredCapabilities.browserVersion,
@@ -97,6 +106,10 @@ export async function launchBrowser(
     existingBrowser.publicAPI.overwriteCommand("deleteSession", async function () {
         await existingBrowser.quit();
         await newBrowser.kill();
+
+        if (filesToRemove.length > 0 && !hasGlobalFilesToRemove()) {
+            await Promise.all(filesToRemove.map(path => fs.remove(path)));
+        }
     });
 
     existingBrowser.publicAPI.addCommand("getDriverPid", () => newBrowser.getDriverPid());

src/browser/standalone/types.ts

@@ -20,6 +20,7 @@ export type StandaloneBrowserOptions = Pick<
     | "user"
     | "key"
     | "system"
+    | "stateOpts"
 >;
 
 export type StandaloneBrowserOptionsInput = Partial<Omit<StandaloneBrowserOptions, "system">> & {

src/browser/types.ts

@@ -9,7 +9,9 @@ import type { Callstack } from "./history/callstack";
 import type { Test, Hook } from "../test-reader/test-object";
 import type { CaptureSnapshotOptions, CaptureSnapshotResult } from "./commands/captureDomSnapshot";
 import type { Options } from "@testplane/wdio-types";
-import type { SaveStateData, SaveStateOptions } from "./commands/saveState";
+import type { SaveStateData } from "./commands/saveState";
+import type { StateOpts } from "../config/types";
+import type { GetStateOptions } from "./commands/getState";
 import type { RestoreStateOptions } from "./commands/restoreState";
 import type { WaitForStaticToLoadResult } from "./commands/waitForStaticToLoad";
 import type { CDP } from "./cdp";
@@ -81,8 +83,9 @@ declare global {
 
             getConfig(this: WebdriverIO.Browser): Promise<BrowserConfig>;
 
-            saveState(options?: SaveStateOptions): Promise<SaveStateData>;
-            restoreState(options: RestoreStateOptions): Promise<void>;
+            saveState(options?: StateOpts): Promise<SaveStateData>;
+            restoreState(options?: RestoreStateOptions): Promise<void>;
+            getState(options?: GetStateOptions): Promise<SaveStateData | null>;
 
             overwriteCommand<CommandName extends BrowserCommand>(
                 name: CommandName,

src/config/browser-options.js

@@ -2,7 +2,7 @@
 
 const _ = require("lodash");
 const fs = require("fs-extra");
-const option = require("gemini-configparser").option;
+const { option, section } = require("gemini-configparser");
 const defaults = require("./defaults");
 const optionsBuilder = require("./options-builder");
 const utils = require("./utils");
@@ -15,7 +15,7 @@ const { extractSelectivityEnabledEnvVariable } = require("./utils");
 const is = utils.is;
 
 function provideRootDefault(name) {
-    return () => defaults[name];
+    return () => _.get(defaults, name);
 }
 
 exports.getTopLevel = () => {
@@ -57,7 +57,7 @@ exports.getPerBrowser = () => {
 
 function provideTopLevelDefault(name) {
     return config => {
-        const value = config[name];
+        const value = _.get(config, name);
 
         if (_.isUndefined(value)) {
             throw new Error(`"${name}" should be set at the top level or per-browser option`);
@@ -443,5 +443,15 @@ function buildBrowserOptions(defaultFactory, extra) {
                 ...extractSelectivityEnabledEnvVariable(ENV_PREFIXES),
             }),
         }),
+
+        stateOpts: section({
+            path: option({
+                defaultValue: defaultFactory("stateOpts.path"),
+            }),
+            cookies: options.optionalBoolean("stateOpts.cookies"),
+            localStorage: options.optionalBoolean("stateOpts.localStorage"),
+            sessionStorage: options.optionalBoolean("stateOpts.sessionStorage"),
+            keepFile: options.optionalBoolean("stateOpts.keepFile"),
+        }),
     });
 }

src/config/defaults.js

@@ -20,6 +20,13 @@ module.exports = {
         clustersSize: 10,
         stopOnFirstFail: false,
     },
+    stateOpts: {
+        path: null,
+        cookies: true,
+        localStorage: true,
+        sessionStorage: true,
+        keepFile: false,
+    },
     buildDiffOpts: {
         ignoreAntialiasing: true,
         ignoreCaret: true,

src/config/types.ts

@@ -1,6 +1,6 @@
 import type { BrowserConfig } from "./browser-config";
 import type { BrowserTestRunEnvOptions } from "../runner/browser-env/vite/types";
-import type { Test } from "../types";
+import type { Cookie, Test } from "../types";
 import type { ChildProcessWithoutNullStreams } from "child_process";
 import type { RequestOptions } from "https";
 import type { Config } from "./index";
@@ -282,6 +282,17 @@ export interface TimeTravelConfig {
     mode: TimeTravelMode;
 }
 
+export type StateOpts = {
+    path?: string;
+
+    cookies?: boolean;
+    localStorage?: boolean;
+    sessionStorage?: boolean;
+
+    cookieFilter?: (cookie: Cookie) => boolean;
+    keepFile?: boolean;
+};
+
 /**
  * @param {Object} dependency - Object with dependency scope and posix relative path
  * @param {"browser"|"testplane"|string} dependency.scope - Dependency scope
@@ -336,6 +347,7 @@ export interface CommonConfig {
     buildDiffOpts: BuildDiffOptsConfig;
     assertViewOpts: AssertViewOpts;
     expectOpts: ExpectOptsConfig;
+    stateOpts?: StateOpts;
     meta: { [name: string]: unknown };
     windowSize: { width: number; height: number } | `${number}x${number}` | null;
     orientation: "landscape" | "portrait" | null;

src/events/index.ts

@@ -44,6 +44,8 @@ export const RunnerSyncEvents = {
 
     DOM_SNAPSHOTS: "domSnapshots",
 
+    ADD_FILE_TO_REMOVE: "addFileToRemove",
+
     TEST_DEPENDENCIES: "testDependencies",
 } as const;