From a5d7842ff61f692537d6a52abde554f6d196f1e2 Mon Sep 17 00:00:00 2001
From: Roman Kuznetsov <kuznetsov.roman.orenburg@gmail.com>
Date: Thu, 27 Jun 2024 15:29:08 +0300
Subject: [PATCH] fix: don't add error snippets on screenshot error

---
 src/error-snippets/index.ts      | 10 +++++-----
 src/error-snippets/utils.ts      | 12 ++++++++++++
 test/src/error-snippets/index.ts | 18 ++++++++++++++++++
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/error-snippets/index.ts b/src/error-snippets/index.ts
index 753012bd1..9cb50ca1d 100644
--- a/src/error-snippets/index.ts
+++ b/src/error-snippets/index.ts
@@ -1,6 +1,6 @@
 import { findRelevantStackFrame, resolveLocationWithStackFrame } from "./frames";
 import { extractSourceMaps, resolveLocationWithSourceMap } from "./source-maps";
-import { getSourceCodeFile, formatErrorSnippet } from "./utils";
+import { getSourceCodeFile, formatErrorSnippet, shouldNotAddCodeSnippet } from "./utils";
 import logger from "../utils/logger";
 import type { ResolvedFrame, SufficientStackFrame, WithSnippetError } from "./types";
 
@@ -14,11 +14,11 @@ const stackFrameLocationResolver = async (stackFrame: SufficientStackFrame): Pro
 };
 
 export const extendWithCodeSnippet = async (err: WithSnippetError): Promise<WithSnippetError> => {
-    if (!err) {
-        return err;
-    }
-
     try {
+        if (shouldNotAddCodeSnippet(err)) {
+            return err;
+        }
+
         const relevantStackFrame = findRelevantStackFrame(err);
 
         if (!relevantStackFrame) {
diff --git a/src/error-snippets/utils.ts b/src/error-snippets/utils.ts
index 45b01e87f..dc8608da2 100644
--- a/src/error-snippets/utils.ts
+++ b/src/error-snippets/utils.ts
@@ -4,6 +4,8 @@ import fs from "fs-extra";
 import { codeFrameColumns } from "@babel/code-frame";
 import { getErrorTitle } from "../browser/stacktrace/utils";
 import { SNIPPET_LINES_ABOVE, SNIPPET_LINES_BELOW, SOURCE_MAP_URL_COMMENT } from "./constants";
+import { AssertViewError } from "../browser/commands/assert-view/errors/assert-view-error";
+import { BaseStateError } from "../browser/commands/assert-view/errors/base-state-error";
 
 interface FormatFileNameHeaderOpts {
     line: number;
@@ -17,6 +19,16 @@ interface FormatErrorSnippetOpts {
     location: { line: number; column: number };
 }
 
+export const shouldNotAddCodeSnippet = (err: Error): boolean => {
+    if (!err) {
+        return true;
+    }
+
+    const isScreenshotError = [AssertViewError, BaseStateError].some(errorClass => err instanceof errorClass);
+
+    return isScreenshotError;
+};
+
 export const softFileURLToPath = (fileName: string): string => {
     if (!fileName.startsWith("file://")) {
         return fileName;
diff --git a/test/src/error-snippets/index.ts b/test/src/error-snippets/index.ts
index a1fdee4eb..78a63dff1 100644
--- a/test/src/error-snippets/index.ts
+++ b/test/src/error-snippets/index.ts
@@ -2,6 +2,9 @@ import fs from "fs-extra";
 import sinon, { type SinonStub } from "sinon";
 import proxyquire from "proxyquire";
 
+import { AssertViewError } from "../../../src/browser/commands/assert-view/errors/assert-view-error";
+import { BaseStateError } from "../../../src/browser/commands/assert-view/errors/base-state-error";
+
 describe("error-snippets", () => {
     const cloneError = (err: Error): Error => {
         const newError = err;
@@ -41,6 +44,21 @@ describe("error-snippets", () => {
 
         afterEach(() => sandbox.restore());
 
+        [AssertViewError, BaseStateError].forEach(ErrorClass => {
+            it(`should not add error snippet for ${ErrorClass.name}`, async () => {
+                const error = new ErrorClass(
+                    "foo",
+                    { path: "/some/path", size: { width: 800, height: 600 } },
+                    { path: "/some/path", relativePath: "../path", size: { width: 800, height: 600 } },
+                );
+                const originalStack = error.stack;
+
+                await extendWithCodeSnippet(error);
+
+                assert.equal(error.stack, originalStack);
+            });
+        });
+
         it("should not modify error, if it is falsy", async () => {
             await assert.eventually.isUndefined(extendWithCodeSnippet());
             await assert.eventually.isNull(extendWithCodeSnippet(null));