Save and restore state

package.json

@@ -32,7 +32,7 @@
     "prettier-watch": "onchange '**' --exclude-path .prettierignore -- prettier --write {{changed}}",
     "test-unit": "_mocha \"test/!(integration)/**/*.js\"",
     "test": "npm run test-unit && npm run check-types && npm run lint",
-    "test-integration": "mocha -r ts-node/register -r test/integration/standalone/preload-browser.fixture.ts test/integration/standalone/standalone.test.ts",
+    "test-integration": "mocha -r ts-node/register -r test/integration/standalone/preload-browser.fixture.ts test/integration/standalone/standalone.test.ts test/integration/standalone/standalone-save-state.test.ts",
     "toc": "doctoc docs --title '### Contents'",
     "precommit": "npm run lint",
     "prepack": "npm run clean && npm run build",
@@ -101,6 +101,7 @@
     "png-validator": "1.1.0",
     "recast": "0.23.6",
     "resolve.exports": "2.0.2",
+    "set-cookie-parser": "2.7.1",
     "sizzle": "2.3.6",
     "socket.io": "4.7.5",
     "socket.io-client": "4.7.5",
@@ -139,6 +140,7 @@
     "@types/mocha": "10.0.1",
     "@types/node": "18.19.3",
     "@types/proxyquire": "1.3.28",
+    "@types/set-cookie-parser": "2.4.10",
     "@types/sinon": "17.0.1",
     "@types/sinonjs__fake-timers": "8.1.2",
     "@types/strftime": "0.9.8",

src/browser/commands/index.ts

@@ -9,4 +9,6 @@ export const customCommandFileNames = [
     "switchToRepl",
     "moveCursorTo",
     "captureDomSnapshot",
+    "saveState",
+    "restoreState",
 ];

src/browser/commands/restoreState.ts

@@ -0,0 +1,4 @@
+import restoreState from "./restoreState/index";
+
+export default restoreState;
+export * from "./restoreState/index";

src/browser/commands/restoreState/clearAllIndexedDB.ts

@@ -0,0 +1,27 @@
+export async function clearAllIndexedDB(): Promise<void> {
+    try {
+        if (!("databases" in indexedDB)) {
+            throw new Error("Your browser don't indexedDB.databases()");
+        }
+
+        const dbList = await indexedDB.databases();
+
+        await Promise.all(
+            dbList.map((dbInfo: { name?: string }): Promise<void> => {
+                if (!dbInfo.name) {
+                    return Promise.resolve();
+                }
+
+                return new Promise<void>((resolve, reject) => {
+                    const deleteReq = indexedDB.deleteDatabase(dbInfo.name as string);
+
+                    deleteReq.addEventListener("success", () => resolve());
+                    deleteReq.addEventListener("error", () => reject(deleteReq.error));
+                    deleteReq.addEventListener("blocked", () => resolve());
+                });
+            }),
+        );
+    } catch (error) {
+        console.error(error);
+    }
+}

src/browser/commands/restoreState/index.ts

@@ -0,0 +1,144 @@
+import fs from "fs-extra";
+import _ from "lodash";
+
+// import { clearAllIndexedDB } from "./clearAllIndexedDB";
+// import { restoreIndexedDB } from "./restoreIndexedDB";
+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, getWebdriverFrames, SaveStateData, SaveStateOptions } from "../saveState";
+import { Protocol } from "devtools-protocol";
+
+export type RestoreStateOptions = SaveStateOptions & {
+    data?: SaveStateData;
+};
+
+export default (browser: Browser): void => {
+    const { publicAPI: session } = browser;
+
+    session.addCommand("restoreState", async (_options: RestoreStateOptions) => {
+        const options = { ...defaultOptions, ..._options };
+
+        let restoreState: SaveStateData | undefined = options.data;
+
+        if (options.path) {
+            restoreState = await fs.readJson(options.path);
+        }
+
+        if (!restoreState) {
+            logger.error("Can't restore state: please provide a path to file or data");
+            return;
+        }
+
+        switch (browser.config.automationProtocol) {
+            case WEBDRIVER_PROTOCOL: {
+                if (restoreState.cookies && options.cookies) {
+                    await session.setCookies(restoreState.cookies);
+                }
+
+                if (restoreState.framesData) {
+                    await session.switchToParentFrame();
+
+                    const frames = await getWebdriverFrames(session);
+
+                    for (let i = -1; i < frames.length; i++) {
+                        await session.switchToParentFrame();
+
+                        // start with -1 for get data from main page
+                        if (i > -1) {
+                            await session.switchFrame(frames[i]);
+                        }
+
+                        const origin = await session.execute<string, []>(() => window.location.origin);
+
+                        const frameData = restoreState.framesData[origin];
+
+                        if (frameData) {
+                            if (frameData.localStorage && options.localStorage) {
+                                await session.execute<void, [Record<string, string>, "localStorage"]>(
+                                    restoreStorage,
+                                    frameData.localStorage,
+                                    "localStorage",
+                                );
+                            }
+
+                            if (frameData.sessionStorage && options.sessionStorage) {
+                                await session.execute<void, [Record<string, string>, "sessionStorage"]>(
+                                    restoreStorage,
+                                    frameData.sessionStorage,
+                                    "sessionStorage",
+                                );
+                            }
+
+                            // @TODO: will make it later
+                            // if (frameData.indexDB && options.indexDB) {
+                            //     // @todo: Doesn't work now
+                            //     await session.execute(clearAllIndexedDB);
+                            //     await session.execute(restoreIndexedDB, frameData.indexDB);
+                            // }
+                        }
+                    }
+
+                    await session.switchToParentFrame();
+                }
+
+                break;
+            }
+            case DEVTOOLS_PROTOCOL: {
+                const puppeteer = await session.getPuppeteer();
+                const pages = await puppeteer.pages();
+                const page = pages[0];
+                const frames = page.frames();
+
+                if (restoreState.cookies && options.cookies) {
+                    await page.setCookie(
+                        ...restoreState.cookies.map(cookie => ({
+                            ...cookie,
+                            sameSite: _.startCase(_.toLower(cookie.sameSite)) as Protocol.Network.CookieSameSite,
+                        })),
+                    );
+                }
+
+                for (const frame of frames) {
+                    const origin = new URL(frame.url()).origin;
+
+                    if (origin === "null" || !restoreState.framesData[origin]) {
+                        continue;
+                    }
+
+                    const frameData = restoreState.framesData[origin];
+
+                    if (!frameData) {
+                        continue;
+                    }
+
+                    if (frameData.localStorage && options.localStorage) {
+                        await frame.evaluate(
+                            restoreStorage,
+                            frameData.localStorage as Record<string, string>,
+                            "localStorage" as const,
+                        );
+                    }
+
+                    if (frameData.sessionStorage && options.sessionStorage) {
+                        await frame.evaluate(
+                            restoreStorage,
+                            frameData.sessionStorage as Record<string, string>,
+                            "sessionStorage" as const,
+                        );
+                    }
+
+                    // @TODO: will make it later
+                    // if (frameData.indexDB) {
+                    //     // @todo: Doesn't work now
+                    //     await frame.evaluate(clearAllIndexedDB);
+                    //     await frame.evaluate(restoreIndexedDB, frameData.indexDB);
+                    // }
+                }
+                break;
+            }
+        }
+    });
+};

src/browser/commands/restoreState/restoreIndexedDB.ts

@@ -0,0 +1,67 @@
+import { DumpIndexDB, DumpStoreIndexDB } from "../saveState/dumpIndexedDB";
+
+export async function restoreIndexedDB(dump: Record<string, DumpIndexDB>): Promise<void> {
+    try {
+        for (const [dbName, dbData] of Object.entries(dump)) {
+            const version = dbData.version || 1;
+            const stores = dbData.stores || {};
+
+            await new Promise<void>((resolve, reject) => {
+                const openReq = indexedDB.open(dbName, version);
+
+                openReq.addEventListener("upgradeneeded", (): void => {
+                    const db = openReq.result;
+
+                    // Restore stores
+                    for (const [storeName, storeInfo] of Object.entries(stores)) {
+                        const { keyPath, autoIncrement, indexes } = storeInfo as DumpStoreIndexDB;
+                        const objectStore = db.createObjectStore(storeName, {
+                            keyPath: keyPath ?? undefined,
+                            autoIncrement: !!autoIncrement,
+                        });
+
+                        for (const idx of indexes) {
+                            objectStore.createIndex(idx.name, idx.keyPath, {
+                                unique: idx.unique,
+                                multiEntry: idx.multiEntry,
+                            });
+                        }
+                    }
+                });
+
+                openReq.addEventListener("success", (): void => {
+                    const db = openReq.result;
+                    const tx = db.transaction(Object.keys(stores), "readwrite");
+
+                    for (const [storeName, storeInfo] of Object.entries(stores)) {
+                        const { records } = storeInfo as DumpStoreIndexDB;
+                        const store = tx.objectStore(storeName);
+
+                        for (const { key, value } of records) {
+                            try {
+                                if (store.keyPath) {
+                                    store.put(value);
+                                } else {
+                                    store.put(value, key as unknown as string);
+                                }
+                            } catch (e) {
+                                console.warn(`Insert error ${storeName}:`, e);
+                            }
+                        }
+                    }
+
+                    tx.addEventListener("success", (): void => {
+                        db.close();
+                        resolve();
+                    });
+
+                    tx.addEventListener("error", (): void => reject(tx.error));
+                });
+
+                openReq.addEventListener("error", (): void => reject(openReq.error));
+            });
+        }
+    } catch (error) {
+        console.error(error);
+    }
+}

src/browser/commands/restoreState/restoreStorage.ts

@@ -0,0 +1,8 @@
+export const restoreStorage = (data: Record<string, string>, type: "localStorage" | "sessionStorage"): void => {
+    const storage = window[type];
+    storage.clear();
+
+    Object.keys(data).forEach(key => {
+        storage.setItem(key, data[key]);
+    });
+};

src/browser/commands/saveState.ts

@@ -0,0 +1,4 @@
+import saveState from "./saveState/index";
+
+export default saveState;
+export * from "./saveState/index";