Skip to content

Commit 9624f8c

Browse files
author
rocketraccoon
committed
fix(testplane): webdriver restore state
1 parent 6245004 commit 9624f8c

File tree

6 files changed

+167
-116
lines changed

6 files changed

+167
-116
lines changed

src/browser/commands/restoreState/index.ts

Lines changed: 100 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -4,103 +4,126 @@ import { clearAllIndexedDB } from "./clearAllIndexedDB";
44
import { restoreIndexedDB } from "./restoreIndexedDB";
55
import { restoreStorage } from "./restoreStorage";
66

7+
import _ from "lodash";
78
import type { Browser } from "../../types";
8-
import { defaultOptions, SaveStateOptions } from "../saveState";
9+
import { DEVTOOLS_PROTOCOL, WEBDRIVER_PROTOCOL } from "../../../constants/config";
10+
import { defaultOptions, getWebdriverFrames, SaveStateData, SaveStateOptions } from "../saveState";
911
import { Protocol } from "devtools-protocol";
10-
import { DumpIndexDB } from "../saveState/dumpIndexedDB";
1112

1213
type RestoreStateOptions = SaveStateOptions;
1314

14-
const normalizeCookies = (cookies: Array<Protocol.Network.CookieParam>): Array<Protocol.Network.CookieParam> =>
15-
cookies
16-
.filter(c => c.name && c.value && c.domain)
17-
.map(c => {
18-
const cookie: Partial<Protocol.Network.CookieParam> = {
19-
name: c.name,
20-
value: c.value,
21-
domain: c.domain,
22-
path: c.path || "/",
23-
};
24-
25-
if (c.expires) {
26-
cookie.expires =
27-
typeof c.expires === "string"
28-
? Math.floor(new Date(c.expires).getTime() / 1000)
29-
: Math.floor(c.expires as number);
30-
}
31-
32-
if (c.secure !== undefined) cookie.secure = c.secure;
33-
if (c.httpOnly !== undefined) cookie.httpOnly = c.httpOnly;
34-
if (c.sameSite) cookie.sameSite = c.sameSite;
35-
36-
if (!c.domain?.startsWith(".")) {
37-
cookie.url = `https://${c.domain}`;
38-
} else {
39-
cookie.url = `https://${c.domain.replace(/^\./, "")}`;
40-
}
41-
42-
return cookie as Protocol.Network.CookieParam;
43-
});
44-
45-
type FrameData = {
46-
localStorage: Record<string, string>;
47-
sessionStorage: Record<string, string>;
48-
indexDB: Record<string, DumpIndexDB>;
49-
};
50-
51-
type RestoreState = {
52-
cookies: Array<Protocol.Network.CookieParam>;
53-
framesData: Record<string, FrameData>;
54-
};
55-
5615
export default (browser: Browser): void => {
5716
const { publicAPI: session } = browser;
5817

5918
session.addCommand("restoreState", async (_options: RestoreStateOptions) => {
6019
const options = { ...defaultOptions, ..._options };
6120

62-
const restoreState: RestoreState = await fs.readJson(options.path);
21+
const restoreState: SaveStateData = await fs.readJson(options.path);
6322

64-
const puppeteer = await session.getPuppeteer();
65-
const pages = await puppeteer.pages();
66-
const page = pages[0];
67-
const frames = page.frames();
23+
switch (browser.config.automationProtocol) {
24+
case WEBDRIVER_PROTOCOL: {
25+
if (restoreState.cookies && options.cookies) {
26+
await session.setCookies(restoreState.cookies);
27+
}
6828

69-
if (restoreState.cookies && options.cookies) {
70-
const normalized = normalizeCookies(restoreState.cookies);
29+
if (restoreState.framesData) {
30+
await session.switchToParentFrame();
7131

72-
await page.setCookie(...normalized);
73-
}
32+
const frames = await getWebdriverFrames(session);
7433

75-
for (const frame of frames) {
76-
const origin = new URL(frame.url()).origin;
34+
for (let i = -1; i < frames.length; i++) {
35+
await session.switchToParentFrame();
7736

78-
if (origin === "null" || !restoreState.framesData[origin]) {
79-
continue;
80-
}
37+
// start with -1 for get data from main page
38+
if (i > -1) {
39+
await session.switchFrame(frames[i]);
40+
}
8141

82-
const frameData = restoreState.framesData[origin];
42+
const origin = await session.execute<string, []>(() => window.location.origin);
8343

84-
if (frameData.localStorage && options.localStorage) {
85-
await frame.evaluate(
86-
restoreStorage,
87-
frameData.localStorage as Record<string, string>,
88-
"localStorage" as const,
89-
);
90-
}
44+
const frameData = restoreState.framesData[origin];
9145

92-
if (frameData.sessionStorage && options.sessionStorage) {
93-
await frame.evaluate(
94-
restoreStorage,
95-
frameData.sessionStorage as Record<string, string>,
96-
"sessionStorage" as const,
97-
);
98-
}
46+
if (frameData) {
47+
if (frameData.localStorage && options.localStorage) {
48+
await session.execute<void, [Record<string, string>, "localStorage"]>(
49+
restoreStorage,
50+
frameData.localStorage,
51+
"localStorage",
52+
);
53+
}
9954

100-
if (frameData.indexDB) {
101-
// @todo: Doesn't work now
102-
await frame.evaluate(clearAllIndexedDB);
103-
await frame.evaluate(restoreIndexedDB, frameData.indexDB);
55+
if (frameData.sessionStorage && options.sessionStorage) {
56+
await session.execute<void, [Record<string, string>, "sessionStorage"]>(
57+
restoreStorage,
58+
frameData.sessionStorage,
59+
"sessionStorage",
60+
);
61+
}
62+
63+
if (frameData.indexDB && options.indexDB) {
64+
// @todo: Doesn't work now
65+
await session.execute(clearAllIndexedDB);
66+
await session.execute(restoreIndexedDB, frameData.indexDB);
67+
}
68+
}
69+
}
70+
71+
await session.switchToParentFrame();
72+
}
73+
74+
break;
75+
}
76+
case DEVTOOLS_PROTOCOL: {
77+
const puppeteer = await session.getPuppeteer();
78+
const pages = await puppeteer.pages();
79+
const page = pages[0];
80+
const frames = page.frames();
81+
82+
if (restoreState.cookies && options.cookies) {
83+
await page.setCookie(
84+
...restoreState.cookies.map(cookie => ({
85+
...cookie,
86+
sameSite: _.startCase(_.toLower(cookie.sameSite)) as Protocol.Network.CookieSameSite,
87+
})),
88+
);
89+
}
90+
91+
for (const frame of frames) {
92+
const origin = new URL(frame.url()).origin;
93+
94+
if (origin === "null" || !restoreState.framesData[origin]) {
95+
continue;
96+
}
97+
98+
const frameData = restoreState.framesData[origin];
99+
100+
if (!frameData) {
101+
continue;
102+
}
103+
104+
if (frameData.localStorage && options.localStorage) {
105+
await frame.evaluate(
106+
restoreStorage,
107+
frameData.localStorage as Record<string, string>,
108+
"localStorage" as const,
109+
);
110+
}
111+
112+
if (frameData.sessionStorage && options.sessionStorage) {
113+
await frame.evaluate(
114+
restoreStorage,
115+
frameData.sessionStorage as Record<string, string>,
116+
"sessionStorage" as const,
117+
);
118+
}
119+
120+
if (frameData.indexDB) {
121+
// @todo: Doesn't work now
122+
await frame.evaluate(clearAllIndexedDB);
123+
await frame.evaluate(restoreIndexedDB, frameData.indexDB);
124+
}
125+
}
126+
break;
104127
}
105128
}
106129
});

src/browser/commands/saveState/dumpIndexedDB.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,14 @@ export interface DumpStoreIndexDB {
1818
}[];
1919
}
2020

21-
export async function dumpIndexedDB(): Promise<Record<string, unknown> | undefined> {
21+
export async function dumpIndexedDB(): Promise<Record<string, DumpIndexDB> | undefined> {
2222
try {
2323
if (!("databases" in indexedDB)) {
2424
throw new Error("Your browser don't indexedDB.databases()");
2525
}
2626

2727
const dbList = await indexedDB.databases(); // список баз
28-
const result: Record<string, unknown> = {};
28+
const result: Record<string, DumpIndexDB> = {};
2929

3030
for (const dbInfo of dbList) {
3131
const name = dbInfo.name;

src/browser/commands/saveState/dumpStorage.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
export type StorageData = {
2-
localStorage?: Record<string, unknown>;
3-
sessionStorage?: Record<string, unknown>;
2+
localStorage?: Record<string, string>;
3+
sessionStorage?: Record<string, string>;
44
};
55

66
export const dumpStorage = (): StorageData => {
7-
const getData = (storage: Storage): Record<string, unknown> | undefined => {
7+
const getData = (storage: Storage): Record<string, string> | undefined => {
88
const data: Record<string, string> = {};
99

1010
for (let i = 0; i < storage.length; i++) {

src/browser/commands/saveState/index.ts

Lines changed: 45 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
import fs from "fs-extra";
22

33
import type { Browser } from "../../types";
4-
import { dumpIndexedDB } from "./dumpIndexedDB";
4+
import { DumpIndexDB, dumpIndexedDB } from "./dumpIndexedDB";
55
import { dumpStorage, StorageData } from "./dumpStorage";
6-
import { Protocol } from "devtools-protocol";
76
import { DEVTOOLS_PROTOCOL, WEBDRIVER_PROTOCOL } from "../../../constants/config";
7+
import { Cookie } from "@testplane/wdio-protocols";
88

99
export type SaveStateOptions = {
1010
path: string;
@@ -15,13 +15,13 @@ export type SaveStateOptions = {
1515
indexDB?: boolean;
1616
};
1717

18-
type FrameData = StorageData & {
19-
indexDB?: Record<string, unknown>;
18+
export type FrameData = StorageData & {
19+
indexDB?: Record<string, DumpIndexDB>;
2020
};
2121

2222
export type SaveStateData = {
23-
cookies?: Array<Protocol.Network.CookieParam>;
24-
framesData?: Record<string, FrameData>;
23+
cookies?: Array<Cookie>;
24+
framesData: Record<string, FrameData>;
2525
};
2626

2727
export const defaultOptions = {
@@ -31,39 +31,49 @@ export const defaultOptions = {
3131
indexDB: false,
3232
};
3333

34+
export const getWebdriverFrames = async (session: WebdriverIO.Browser): Promise<string[]> =>
35+
session.execute<string[], []>(() =>
36+
Array.from(document.getElementsByTagName("iframe"))
37+
.map(el => el.getAttribute("src") as string)
38+
.filter(src => src !== null && src !== "about:blank"),
39+
);
40+
3441
export default (browser: Browser): void => {
3542
const { publicAPI: session } = browser;
3643

3744
session.addCommand("saveState", async (_options: SaveStateOptions) => {
3845
const options = { ...defaultOptions, ..._options };
3946

40-
const data: SaveStateData = {};
47+
const data: SaveStateData = {
48+
framesData: {},
49+
};
4150

4251
switch (browser.config.automationProtocol) {
4352
case WEBDRIVER_PROTOCOL: {
4453
if (options.cookies) {
4554
const storageCookies = await session.storageGetCookies({});
4655

4756
data.cookies = storageCookies.cookies.map(cookie => ({
48-
...cookie,
57+
name: cookie.name,
4958
value: cookie.value.value,
50-
sameSite: cookie.sameSite.toLowerCase() as Protocol.Network.CookieSameSite,
59+
domain: cookie.domain,
60+
path: cookie.path,
61+
expires: cookie.expiry,
62+
httpOnly: cookie.httpOnly,
63+
secure: cookie.secure,
64+
sameSite: cookie.sameSite,
5165
}));
5266
}
5367

5468
await session.switchToParentFrame();
5569

56-
const frames = await session.execute<string[], []>(() =>
57-
Array.from(document.getElementsByTagName("iframe"))
58-
.map(el => el.getAttribute("src") as string)
59-
.filter(src => src !== null && src !== "about:blank"),
60-
);
61-
70+
const frames = await getWebdriverFrames(session);
6271
const framesData: Record<string, FrameData> = {};
6372

6473
for (let i = -1; i < frames.length; i++) {
6574
await session.switchToParentFrame();
6675

76+
// start with -1 for get data from main page
6777
if (i > -1) {
6878
await session.switchFrame(frames[i]);
6979
}
@@ -74,20 +84,22 @@ export default (browser: Browser): void => {
7484
continue;
7585
}
7686

77-
const { localStorage, sessionStorage } = await session.execute<StorageData, []>(dumpStorage);
78-
7987
const frameData: FrameData = {};
8088

81-
if (localStorage && options.localStorage) {
82-
frameData.localStorage = localStorage;
83-
}
89+
if (options.localStorage || options.sessionStorage) {
90+
const { localStorage, sessionStorage } = await session.execute<StorageData, []>(dumpStorage);
91+
92+
if (localStorage && options.localStorage) {
93+
frameData.localStorage = localStorage;
94+
}
8495

85-
if (sessionStorage && options.sessionStorage) {
86-
frameData.sessionStorage = sessionStorage;
96+
if (sessionStorage && options.sessionStorage) {
97+
frameData.sessionStorage = sessionStorage;
98+
}
8799
}
88100

89101
if (options.indexDB) {
90-
const indexDB: Record<string, unknown> | undefined = await session.execute(dumpIndexedDB);
102+
const indexDB: Record<string, DumpIndexDB> | undefined = await session.execute(dumpIndexedDB);
91103

92104
if (indexDB) {
93105
frameData.indexDB = indexDB;
@@ -120,20 +132,22 @@ export default (browser: Browser): void => {
120132
continue;
121133
}
122134

123-
const { localStorage, sessionStorage }: StorageData = await frame.evaluate(dumpStorage);
124-
125135
const frameData: FrameData = {};
126136

127-
if (localStorage && options.localStorage) {
128-
frameData.localStorage = localStorage;
129-
}
137+
if (options.localStorage || options.sessionStorage) {
138+
const { localStorage, sessionStorage }: StorageData = await frame.evaluate(dumpStorage);
130139

131-
if (sessionStorage && options.sessionStorage) {
132-
frameData.sessionStorage = sessionStorage;
140+
if (localStorage && options.localStorage) {
141+
frameData.localStorage = localStorage;
142+
}
143+
144+
if (sessionStorage && options.sessionStorage) {
145+
frameData.sessionStorage = sessionStorage;
146+
}
133147
}
134148

135149
if (options.indexDB) {
136-
const indexDB: Record<string, unknown> | undefined = await frame.evaluate(dumpIndexedDB);
150+
const indexDB: Record<string, DumpIndexDB> | undefined = await frame.evaluate(dumpIndexedDB);
137151

138152
if (indexDB) {
139153
frameData.indexDB = indexDB;

0 commit comments

Comments
 (0)