Save

code/lib/instrumenter/src/instrumenter.ts

@@ -49,7 +49,18 @@ const isInstrumentable = (o: unknown) => {
   if (o.constructor === undefined) return true;
   const proto = o.constructor.prototype;
   if (!isObject(proto)) return false;
-  if (Object.prototype.hasOwnProperty.call(proto, 'isPrototypeOf') === false) return false;
+  // if (Object.prototype.hasOwnProperty.call(proto, 'isPrototypeOf') === false) return false;
+  return true;
+};
+
+const old = (o: unknown) => {
+  if (isObject(o) || isModule(o)) return true;
+  return false;
+
+  if (o.constructor === undefined) return true;
+  const proto = o.constructor.prototype;
+  if (!isObject(proto)) return false;
+  // if (Object.prototype.hasOwnProperty.call(proto, 'isPrototypeOf') === false) return false;
   return true;
 };
 
@@ -294,7 +305,9 @@ export class Instrumenter {
     if (!isInstrumentable(obj)) return obj;
 
     const { mutate = false, path = [] } = options;
-    return Object.keys(obj).reduce(
+
+    const keys = options.getKeys ? options.getKeys(obj) : Object.keys(obj);
+    return keys.reduce(
       (acc, key) => {
         const value = (obj as Record<string, any>)[key];
 
@@ -305,13 +318,13 @@ export class Instrumenter {
         }
 
         // Already patched, so we pass through unchanged
-        if (typeof value.__originalFn__ === 'function') {
+        if ('__originalFn__' in value && typeof value.__originalFn__ === 'function') {
           acc[key] = value;
           return acc;
         }
 
         // Patch the function and mark it "patched" by adding a reference to the original function
-        acc[key] = (...args: any[]) => this.track(key, value, args, options);
+        acc[key] = (...args: any[]) => this.track(key, value, acc, args, options);
         acc[key].__originalFn__ = value;
 
         // Reuse the original name as the patched function's name
@@ -334,7 +347,7 @@ export class Instrumenter {
   // Monkey patch an object method to record calls.
   // Returns a function that invokes the original function, records the invocation ("call") and
   // returns the original result.
-  track(method: string, fn: Function, args: any[], options: Options) {
+  track(method: string, fn: Function, object: object, args: any[], options: Options) {
     const storyId: StoryId =
       args?.[0]?.__storyId__ || global.__STORYBOOK_PREVIEW__?.selectionStore?.selection?.storyId;
     const { cursor, ancestors } = this.getState(storyId);
@@ -344,11 +357,11 @@ export class Instrumenter {
     const interceptable = typeof intercept === 'function' ? intercept(method, path) : intercept;
     const call = { id, cursor, storyId, ancestors, path, method, args, interceptable, retain };
     const interceptOrInvoke = interceptable && !ancestors.length ? this.intercept : this.invoke;
-    const result = interceptOrInvoke.call(this, fn, call, options);
+    const result = interceptOrInvoke.call(this, fn, object, call, options);
     return this.instrument(result, { ...options, mutate: true, path: [{ __callId__: call.id }] });
   }
 
-  intercept(fn: Function, call: Call, options: Options) {
+  intercept(fn: Function, object: object, call: Call, options: Options) {
     const { chainedCallIds, isDebugging, playUntil } = this.getState(call.storyId);
 
     // For a "jump to step" action, continue playing until we hit a call by that ID.
@@ -358,7 +371,7 @@ export class Instrumenter {
       if (playUntil === call.id) {
         this.setState(call.storyId, { playUntil: undefined });
       }
-      return this.invoke(fn, call, options);
+      return this.invoke(fn, object, call, options);
     }
 
     // Instead of invoking the function, defer the function call until we continue playing.
@@ -373,11 +386,11 @@ export class Instrumenter {
         const { [call.id]: _, ...resolvers } = state.resolvers;
         return { isLocked: true, resolvers };
       });
-      return this.invoke(fn, call, options);
+      return this.invoke(fn, object, call, options);
     });
   }
 
-  invoke(fn: Function, call: Call, options: Options) {
+  invoke(fn: Function, object: object, call: Call, options: Options) {
     // TODO this doesnt work because the abortSignal we have here is the newly created one
     // const { abortSignal } = global.window.__STORYBOOK_PREVIEW__ || {};
     // if (abortSignal && abortSignal.aborted) throw IGNORED_EXCEPTION;
@@ -510,7 +523,7 @@ export class Instrumenter {
         };
       });
 
-      const result = fn(...finalArgs);
+      const result = fn.apply(object, finalArgs);
 
       // Track the result so we can trace later uses of it back to the originating call.
       // Primitive results (undefined, null, boolean, string, number, BigInt) are ignored.

code/lib/instrumenter/src/types.ts

@@ -90,4 +90,5 @@ export interface Options {
   mutate?: boolean;
   path?: Array<string | CallRef>;
   getArgs?: (call: Call, state: State) => Call['args'];
+  getKeys?: (obj: unknown) => string[];
 }

code/lib/test/src/expect.ts

@@ -0,0 +1,84 @@
+import * as chai from 'chai';
+import { GLOBAL_EXPECT, getState, setState } from '@vitest/expect';
+import type { Assertion, ExpectStatic } from '@vitest/expect';
+import type { MatcherState } from '@vitest/expect';
+
+export function createExpect(test?: any) {
+  const expect: ExpectStatic = chai.expect as ExpectStatic;
+
+  expect.getState = () => getState<MatcherState>(expect);
+  expect.setState = (state) => setState(state as Partial<MatcherState>, expect);
+
+  // @ts-expect-error global is not typed
+  const globalState = getState(globalThis[GLOBAL_EXPECT]) || {};
+
+  setState<MatcherState>(
+    {
+      // this should also add "snapshotState" that is added conditionally
+      ...globalState,
+      assertionCalls: 0,
+      isExpectingAssertions: false,
+      isExpectingAssertionsError: null,
+      expectedAssertionsNumber: null,
+      expectedAssertionsNumberErrorGen: null,
+      testPath: test ? test.suite.file?.filepath : globalState.testPath,
+    },
+    expect
+  );
+
+  // @ts-expect-error untyped
+  expect.extend = (matchers) => chai.expect.extend(expect, matchers);
+
+  expect.soft = (...args) => {
+    const assert = expect(...args);
+    expect.setState({
+      soft: true,
+    });
+    return assert;
+  };
+
+  expect.unreachable = (message?: string) => {
+    chai.assert.fail(`expected${message ? ` "${message}" ` : ' '}not to be reached`);
+  };
+
+  function assertions(expected: number) {
+    const errorGen = () =>
+      new Error(
+        `expected number of assertions to be ${expected}, but got ${
+          expect.getState().assertionCalls
+        }`
+      );
+    if (Error.captureStackTrace) Error.captureStackTrace(errorGen(), assertions);
+
+    expect.setState({
+      expectedAssertionsNumber: expected,
+      expectedAssertionsNumberErrorGen: errorGen,
+    });
+  }
+
+  function hasAssertions() {
+    const error = new Error('expected any number of assertion, but got none');
+    if (Error.captureStackTrace) Error.captureStackTrace(error, hasAssertions);
+
+    expect.setState({
+      isExpectingAssertions: true,
+      isExpectingAssertionsError: error,
+    });
+  }
+
+  // chai.util.addMethod(expect, 'assertions', assertions);
+  // chai.util.addMethod(expect, 'hasAssertions', hasAssertions);
+
+  return expect;
+}
+
+const globalExpect = createExpect();
+
+Object.defineProperty(globalThis, GLOBAL_EXPECT, {
+  value: globalExpect,
+  writable: true,
+  configurable: true,
+});
+
+export { assert, should } from 'chai';
+export { chai, globalExpect as expect };

code/lib/test/src/index.ts

@@ -1,11 +1,23 @@
 /* eslint-disable import/no-extraneous-dependencies,import/no-named-default */
-import { default as expectPatched } from '@storybook/expect';
 import { instrument } from '@storybook/instrumenter';
 import * as matchers from '@testing-library/jest-dom/matchers';
-import * as _mock from '@vitest/spy';
+import * as spy from '@vitest/spy';
+import chai from 'chai';
+import { JestAsymmetricMatchers, JestChaiExpect, JestExtend } from '@vitest/expect';
+import { default as expectPatched } from '@storybook/expect';
+
+import { expect as _expect } from './expect';
+
+chai.use(JestExtend);
+chai.use(JestChaiExpect);
+chai.use(JestAsymmetricMatchers);
+// expect.extend(matchers);
+
 export type * from '@vitest/spy';
+export * from '@vitest/spy';
+// export { expect };
 
-export const { mock } = instrument({ mock: _mock }, { retain: true });
+export const { fn } = instrument({ fn: spy.fn }, { retain: true });
 
 /**
  * The `expect` function is used every time you want to test a value.
@@ -23,10 +35,40 @@ export interface Expect extends Pick<jest.Expect, keyof jest.Expect> {
     jest.Matchers<Promise<void>, T>
   >;
 }
+//
+// const expect2 = (value) => {
+//   return {
+//     toBe: (other) => {
+//       throw new Error('bla');
+//     },
+//   };
+// };
+//
+// Object.assign(expect2, {
+//   bla: () => {},
+// });
 
-expectPatched.extend(matchers);
+// export const jestExpect: Expect = instrument(
+//   { expect: expectPatched },
+//   { intercept: (_method, path) => path[0] !== 'expect' }
+// ).expect as unknown as Expect;
 
 export const expect: Expect = instrument(
-  { expect: expectPatched },
-  { intercept: (_method, path) => path[0] !== 'expect' }
+  { expect: _expect },
+  {
+    getKeys: (obj) => {
+      if (Object.keys(obj).filter((it) => !it.startsWith('__')).length === 0) {
+        let allKeys = [...Object.keys(Object.getPrototypeOf(obj)), ...Object.keys(obj)].filter(
+          (it) => it !== 'assert' && it !== '__methods'
+        );
+        console.log(allKeys);
+        return allKeys;
+      } else {
+        return Object.keys(obj);
+      }
+    },
+    intercept: (method, path) => {
+      return true;
+    },
+  }
 ).expect as unknown as Expect;