feat: add ability to capture DOM snapshots of specific elements (#1098)

* feat: add ability to capture DOM snapshots of specific elements

* fix: fix review issues regarding undefined usage

Author: shadowusr

src/browser/commands/captureDomSnapshot.ts

@@ -16,14 +16,37 @@ export interface CaptureSnapshotResult {
     textWasTruncated: boolean;
 }
 
-export const captureDomSnapshotInBrowser = ({
-    includeTags = [],
-    includeAttrs = [],
-    excludeTags = [],
-    excludeAttrs = [],
-    truncateText = true,
-    maxTextLength = 100,
-}: CaptureSnapshotOptions = {}): CaptureSnapshotResult => {
+export const captureDomSnapshotInBrowser = (
+    selectorOrElementOrOptions?: string | WebdriverIO.Element | CaptureSnapshotOptions,
+    maybeOptions?: CaptureSnapshotOptions,
+): CaptureSnapshotResult => {
+    let selector: string | null = null;
+    let element: Element | null = null;
+    let options: CaptureSnapshotOptions;
+
+    if (typeof selectorOrElementOrOptions === "string") {
+        selector = selectorOrElementOrOptions;
+        options = maybeOptions || {};
+    } else if (
+        selectorOrElementOrOptions &&
+        typeof selectorOrElementOrOptions === "object" &&
+        "tagName" in selectorOrElementOrOptions
+    ) {
+        element = selectorOrElementOrOptions as Element;
+        options = maybeOptions || {};
+    } else {
+        options = (selectorOrElementOrOptions as CaptureSnapshotOptions) || {};
+    }
+
+    const {
+        includeTags = [],
+        includeAttrs = [],
+        excludeTags = [],
+        excludeAttrs = [],
+        truncateText = true,
+        maxTextLength = 100,
+    } = options;
+
     const BASE_EXCLUDED_TAGS = new Set([
         "HEAD",
         "LINK",
@@ -403,14 +426,39 @@ export const captureDomSnapshotInBrowser = ({
         }
     }
 
-    const startElement = document.body || document.documentElement;
-    if (!startElement) {
-        return {
-            snapshot: "# No elements found",
-            omittedTags: [],
-            omittedAttributes: [],
-            textWasTruncated: false,
-        };
+    let startElement: Element | null;
+
+    if (element) {
+        startElement = element;
+    } else if (selector) {
+        try {
+            startElement = document.querySelector(selector);
+            if (!startElement) {
+                return {
+                    snapshot: `# Element not found: ${selector}`,
+                    omittedTags: [],
+                    omittedAttributes: [],
+                    textWasTruncated: false,
+                };
+            }
+        } catch (error) {
+            return {
+                snapshot: `# Invalid selector: ${selector}`,
+                omittedTags: [],
+                omittedAttributes: [],
+                textWasTruncated: false,
+            };
+        }
+    } else {
+        startElement = document.body || document.documentElement;
+        if (!startElement) {
+            return {
+                snapshot: "# No elements found",
+                omittedTags: [],
+                omittedAttributes: [],
+                textWasTruncated: false,
+            };
+        }
     }
 
     const rootNode = buildElementNode(startElement);
@@ -430,8 +478,23 @@ export default (browser: Browser): void => {
 
     session.addCommand(
         "unstable_captureDomSnapshot",
-        async function (options: Partial<CaptureSnapshotOptions> = {}): Promise<CaptureSnapshotResult> {
-            return session.execute(captureDomSnapshotInBrowser, options);
+        async function (
+            this: WebdriverIO.Browser,
+            selectorOrOptions?: string | Partial<CaptureSnapshotOptions>,
+            maybeOptions?: Partial<CaptureSnapshotOptions>,
+        ): Promise<CaptureSnapshotResult> {
+            return session.execute(captureDomSnapshotInBrowser, selectorOrOptions, maybeOptions);
+        },
+    );
+
+    session.addCommand(
+        "unstable_captureDomSnapshot",
+        async function (
+            this: WebdriverIO.Element,
+            options: Partial<CaptureSnapshotOptions> = {},
+        ): Promise<CaptureSnapshotResult> {
+            return this.execute(captureDomSnapshotInBrowser, options);
         },
+        true,
     );
 };

src/browser/types.ts

@@ -153,10 +153,15 @@ declare global {
 
             clearSession: (this: WebdriverIO.Browser) => Promise<void>;
 
-            unstable_captureDomSnapshot: (
+            unstable_captureDomSnapshot(
                 this: WebdriverIO.Browser,
-                options: CaptureSnapshotOptions,
-            ) => Promise<CaptureSnapshotResult>;
+                options?: Partial<CaptureSnapshotOptions>,
+            ): Promise<CaptureSnapshotResult>;
+            unstable_captureDomSnapshot(
+                this: WebdriverIO.Browser,
+                selector: string,
+                options?: Partial<CaptureSnapshotOptions>,
+            ): Promise<CaptureSnapshotResult>;
         }
 
         interface Element {
@@ -193,6 +198,11 @@ declare global {
             assertView: AssertViewElementCommand;
 
             moveCursorTo: MoveCursorToCommand;
+
+            unstable_captureDomSnapshot(
+                this: WebdriverIO.Element,
+                options?: Partial<CaptureSnapshotOptions>,
+            ): Promise<CaptureSnapshotResult>;
         }
     }
 }

test/src/browser/commands/captureDomSnapshot.ts

@@ -689,4 +689,97 @@ describe('"captureDomSnapshot" command', () => {
             assert.deepEqual(result.snapshot.split("\n"), expected.split("\n"));
         });
     });
+
+    describe("selector-based capture", () => {
+        beforeEach(() => {
+            document.body.innerHTML = `
+                <div class="container">
+                    <h1>Header</h1>
+                    <div class="content">
+                        <p>Paragraph content</p>
+                        <button id="test-button">Click me</button>
+                    </div>
+                    <footer class="footer">Footer content</footer>
+                </div>
+            `;
+        });
+
+        it("should capture snapshot of specific element by selector", () => {
+            const result = captureDomSnapshotInBrowser(".content");
+
+            assert.match(result.snapshot, /div\.content.*@hidden.*:/);
+            assert.match(result.snapshot, /p.*@hidden.*"Paragraph content"/);
+            assert.match(result.snapshot, /button#test-button.*@hidden.*"Click me"/);
+            assert.notMatch(result.snapshot, /h1.*"Header"/);
+            assert.notMatch(result.snapshot, /footer.*"Footer content"/);
+        });
+
+        it("should return error message for non-existent selector", () => {
+            const result = captureDomSnapshotInBrowser(".non-existent");
+
+            assert.equal(result.snapshot, "# Element not found: .non-existent");
+            assert.deepEqual(result.omittedTags, []);
+            assert.deepEqual(result.omittedAttributes, []);
+            assert.equal(result.textWasTruncated, false);
+        });
+
+        it("should return error message for invalid selector", () => {
+            const result = captureDomSnapshotInBrowser("[invalid selector");
+
+            assert.equal(result.snapshot, "# Invalid selector: [invalid selector");
+            assert.deepEqual(result.omittedTags, []);
+            assert.deepEqual(result.omittedAttributes, []);
+            assert.equal(result.textWasTruncated, false);
+        });
+
+        it("should support options with selector", () => {
+            const result = captureDomSnapshotInBrowser(".content", {
+                truncateText: false,
+                maxTextLength: 5,
+            });
+
+            assert.match(result.snapshot, /div\.content.*@hidden.*:/);
+            assert.match(result.snapshot, /p.*@hidden.*"Paragraph content"/);
+            assert.equal(result.textWasTruncated, false);
+        });
+    });
+
+    describe("element-based capture", () => {
+        let testElement: Element;
+
+        beforeEach(() => {
+            document.body.innerHTML = `
+                <div class="container">
+                    <h1>Header</h1>
+                    <div class="content" id="target-content">
+                        <p>Paragraph content</p>
+                        <button id="test-button">Click me</button>
+                    </div>
+                    <footer class="footer">Footer content</footer>
+                </div>
+            `;
+            testElement = document.querySelector(".content") as Element;
+        });
+
+        it("should capture snapshot of specific element passed directly", () => {
+            const result = captureDomSnapshotInBrowser(testElement as any);
+
+            assert.match(result.snapshot, /div\.content#target-content.*@hidden.*:/);
+            assert.match(result.snapshot, /p.*@hidden.*"Paragraph content"/);
+            assert.match(result.snapshot, /button#test-button.*@hidden.*"Click me"/);
+            assert.notMatch(result.snapshot, /h1.*"Header"/);
+            assert.notMatch(result.snapshot, /footer.*"Footer content"/);
+        });
+
+        it("should support options with element", () => {
+            const result = captureDomSnapshotInBrowser(testElement as any, {
+                maxTextLength: 5,
+                truncateText: true,
+            });
+
+            assert.match(result.snapshot, /div\.content#target-content.*@hidden.*:/);
+            assert.match(result.snapshot, /p.*@hidden.*"Parag\.\.\."/);
+            assert.equal(result.textWasTruncated, true);
+        });
+    });
 });