From 2ca4a01af2afa2945f9e85db309c656090f48b80 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Thu, 29 Aug 2024 10:16:54 +0200
Subject: [PATCH] feat(test runner): `expect.extendImmutable`

---
 packages/playwright/src/matchers/expect.ts | 68 +++++++++++++++-------
 tests/playwright-test/expect.spec.ts       | 41 +++++++++++++
 2 files changed, 89 insertions(+), 20 deletions(-)

diff --git a/packages/playwright/src/matchers/expect.ts b/packages/playwright/src/matchers/expect.ts
index ea796bfc72..6bd52b9265 100644
--- a/packages/playwright/src/matchers/expect.ts
+++ b/packages/playwright/src/matchers/expect.ts
@@ -61,6 +61,7 @@ import {
 import { zones } from 'playwright-core/lib/utils';
 import { TestInfoImpl } from '../worker/testInfo';
 import { ExpectError, isExpectError } from './matcherHint';
+import { randomUUID } from 'node:crypto';
 
 // #region
 // Mirrored from https://github.com/facebook/jest/blob/f13abff8df9a0e1148baf3584bcde6d1b479edc7/packages/expect/src/print.ts
@@ -104,11 +105,30 @@ export const printReceivedStringContainExpectedResult = (
 
 type ExpectMessage = string | { message?: string };
 
-function createMatchers(actual: unknown, info: ExpectMetaInfo): any {
-  return new Proxy(expectLibrary(actual), new ExpectMetaInfoProxyHandler(info));
+function createMatchers(actual: unknown, info: ExpectMetaInfo, prefix: string[]): any {
+  return new Proxy(expectLibrary(actual), new ExpectMetaInfoProxyHandler(info, prefix));
 }
 
-function createExpect(info: ExpectMetaInfo) {
+function createExpect(info: ExpectMetaInfo, prefix: string[] = []) {
+
+  function extend(matchers: any, qualifier: string) {
+    const wrappedMatchers: any = {};
+    for (const [name, matcher] of Object.entries(matchers)) {
+      wrappedMatchers[qualifier + name] = function(...args: any[]) {
+        const { isNot, promise, utils } = this;
+        const newThis: ExpectMatcherState = {
+          isNot,
+          promise,
+          utils,
+          timeout: currentExpectTimeout()
+        };
+        (newThis as any).equals = throwUnsupportedExpectMatcherError;
+        return (matcher as any).call(newThis, ...args);
+      };
+    }
+    expectLibrary.extend(wrappedMatchers);
+  }
+
   const expectInstance: Expect<{}> = new Proxy(expectLibrary, {
     apply: function(target: any, thisArg: any, argumentsList: [unknown, ExpectMessage?]) {
       const [actual, messageOrOptions] = argumentsList;
@@ -119,7 +139,7 @@ function createExpect(info: ExpectMetaInfo) {
           throw new Error('`expect.poll()` accepts only function as a first argument');
         newInfo.generator = actual as any;
       }
-      return createMatchers(actual, newInfo);
+      return createMatchers(actual, newInfo, prefix);
     },
 
     get: function(target: any, property: string) {
@@ -128,25 +148,20 @@ function createExpect(info: ExpectMetaInfo) {
 
       if (property === 'extend') {
         return (matchers: any) => {
-          const wrappedMatchers: any = {};
-          for (const [name, matcher] of Object.entries(matchers)) {
-            wrappedMatchers[name] = function(...args: any[]) {
-              const { isNot, promise, utils } = this;
-              const newThis: ExpectMatcherState = {
-                isNot,
-                promise,
-                utils,
-                timeout: currentExpectTimeout()
-              };
-              (newThis as any).equals = throwUnsupportedExpectMatcherError;
-              return (matcher as any).call(newThis, ...args);
-            };
-          }
-          expectLibrary.extend(wrappedMatchers);
+          extend(matchers, '');
           return expectInstance;
         };
       }
 
+      if (property === 'extendImmutable') {
+        return (matchers: any) => {
+          const key = randomUUID();
+          const qualifier = [...prefix, key];
+          extend(matchers, `${qualifier.join(':')}$`);
+          return createExpect(info, qualifier);
+        };
+      }
+
       if (property === 'soft') {
         return (actual: unknown, messageOrOptions?: ExpectMessage) => {
           return configure({ soft: true })(actual, messageOrOptions) as any;
@@ -241,13 +256,26 @@ type ExpectMetaInfo = {
 
 class ExpectMetaInfoProxyHandler implements ProxyHandler<any> {
   private _info: ExpectMetaInfo;
+  private _prefix: string[];
 
-  constructor(info: ExpectMetaInfo) {
+  constructor(info: ExpectMetaInfo, prefix: string[]) {
     this._info = { ...info };
+    this._prefix = prefix;
   }
 
   get(target: Object, matcherName: string | symbol, receiver: any): any {
     let matcher = Reflect.get(target, matcherName, receiver);
+
+    if (typeof matcherName === 'string') {
+      for (let i = this._prefix.length; i > 0; i--) {
+        const qualifiedName = `${this._prefix.slice(0, i).join(':')}$${matcherName}`;
+        if (Reflect.has(target, qualifiedName)) {
+          matcher = Reflect.get(target, qualifiedName, receiver);
+          break;
+        }
+      }
+    }
+
     if (typeof matcherName !== 'string')
       return matcher;
     if (matcher === undefined)
diff --git a/tests/playwright-test/expect.spec.ts b/tests/playwright-test/expect.spec.ts
index 63928e86fb..1f9bd72141 100644
--- a/tests/playwright-test/expect.spec.ts
+++ b/tests/playwright-test/expect.spec.ts
@@ -1063,3 +1063,44 @@ test('should throw error when using .equals()', async ({ runInlineTest }) => {
   expect(result.exitCode).toBe(0);
   expect(result.passed).toBe(1);
 });
+
+test('expect.extendImmutable should work', async ({ runInlineTest }) => {
+  const result = await runInlineTest({
+    'expect-test.spec.ts': `
+      import { test, expect } from '@playwright/test';
+      const expectFoo = expect.extendImmutable({
+        toFoo() {
+          console.log('%%foo');
+          return { pass: true };
+        }
+      });
+      const expectFoo2 = expect.extendImmutable({
+        toFoo() {
+          console.log('%%foo2');
+          return { pass: true };
+        }
+      });
+      const expectBar = expectFoo.extendImmutable({
+        toBar() {
+          console.log('%%bar');
+          return { pass: true };
+        }
+      });
+      test('logs', () => {
+        expect(expectFoo).not.toBe(expectFoo2);
+        expect(expectFoo).not.toBe(expectBar);
+
+        expectFoo().toFoo();
+        expectFoo2().toFoo();
+        expectBar().toFoo();
+        expectBar().toBar();
+      });
+    `
+  });
+  expect(result.outputLines).toEqual([
+    'foo',
+    'foo2',
+    'foo',
+    'bar',
+  ]);
+});
\ No newline at end of file