fix: workaround bug in vitest where getMockImplementation() does not …

[tommie-lie]

…return the implementation used when calling after resetting the mock

This is maybe a bit exotic and will be obsoleted with vitest@4 anyway, but here goes:

Steps to reproduce

const spy = vi.fn((n) => 2 * n)
vi.resetAllMocks()
subject.when(spy).calledWith(1).thenReturn(4)
expect(spy(2)).toEqual(4)

Expected behavior

Test succeeds

Actual behavior

Test fails because spy(2) returns undefined

Explanation and analysis

We have a mock which is setup during the setup-tests phase with a default implementation (vi.fn(impl)) but we have also enabled vitest's restoreMocks setting enabled (which basically calls vi.restoreAllMocks() before each test).
When calling those mocks, we always get the default implementation as state in vitest's docs: https://vitest.dev/api/mock.html#mockrestore

vitest@3 creates a wrapper around tinyspy and keeps its own copy of the mock implementation. When resetting, this wrapped field is reset but the tinyspy keeps its implementation. vitest-when uses getMockImplementation() which returns the wrapper's field, which was reset. In reality, calling the mock will call tinyspy's internal mock implementation.

In vitest@4, they've replaced the wrapper with their own implementation of a spy/mock which behaves correctly and returns the original mock implementation even after resetting the mock.

This PR works around this bug for everyone stuck on vitest@3 but with some rather questionable tricks to be compatible with vitest@3 and vitest@4, ie. with tinyspy installed and not. I could understand if you don't merge it :-)

[mcous]

This is a fun issue, thanks for digging into it! I think a workaround is appropriate - I can see this being really annoying and I'm a heavy user of both Vitest 3 and resetAllMocks

I think I'd like to take a different approach to extracting the original implementation from tinyspy to avoid TLA and module side effects, and I want to make sure any hackiness is isolated and explained, but otherwise I support this!

[mcous]

suggestion: I'm not above adding a bit of hackiness around this for the sake of Vitest 3, but I'm unwilling to add a top-level await / module import side effects to make it happen.

As long as we're hacking, I think there's a way to do this synchronously using Object.getOwnPropertySymbols:

for (const key of Object.getOwnPropertySymbols(spy)) {
  const maybeSpyInternals = (spy as unknown as { [key]: unknown })[key]
  if (
    maybeSpyInternals &&
    typeof maybeSpyInternals === 'object' &&
    'getOriginal' in maybeSpyInternals &&
    typeof maybeSpyInternals.getOriginal === 'function'
  ) {
    // eslint-disable-next-line @typescript-eslint/no-unsafe-call
    return maybeSpyInternals.getOriginal() as TFunc
  }
}

Structurally, I'd want this hackiness to be pretty isolated from the rest of the code (where we can add sufficient comments explaining ourselves and our crimes), so maybe we create a fallback-implementation.ts module or similar to wrap both this and/or the getMockImplementation call. That way, stubs.ts ends up looking something like:

const behaviors = createBehaviorStack<TFunc>()
const fallbackImplementation = getFallbackImplementation(spy)

[tommie-lie]

My first experimental version used Symbol.for("tinyspy:spy") on the spy object until I noticed that getInternalState is actually a public function. But I understand wanting to avoid top-level dynamic imports. We've had our fair share of problems with that in vitest as well and your approach is definitely better.

I guess using the known symbol name will maybe avoid false positives in case some other weird thing adds a getOriginal function. On the other hand, it breaks compatible drop-in replacements for tinyspy. Both events are equally unlikely, though.

Would you consider this a better approach?:

  interface TinyspyInternalState<TFunc extends AnyCallable> {
    getOriginal: () => TFunc
  }

  for (const key of Object.getOwnPropertySymbols(maybeTinyspy)) {
    if (Symbol.keyFor(key) === 'tinyspy:spy') {
      // @ts-expect-error The key was retrieved from maybeTinyspy, so it exists
      const tinySpyInternals = maybeTinyspy[key] as TinyspyInternalState<TFunc>
      return tinySpyInternals.getOriginal()
    }
  }

I could even make the predicate a type predicate and avoid the typecast altogether.

[mcous]

My gut says to go with the duck-typing approach over checking (or even caring at all about) the symbol name. That way, we get a "guarantee" that we're not going to cause a runtime exception, and we only have to work around the type-checker a little bit.

The type-checker is even easier to work around if we make this little addition to the MockInstance type:

  export interface MockInstance<TFunc extends AnyCallable = AnyCallable> {
    getMockName(): string
    getMockImplementation(): TFunc | undefined
    mockImplementation: (impl: TFunc) => this
    mock: MockContext<TFunc>
+   [key: PropertyKey]: unknown
  }

[tommie-lie]

I updated the PR with your suggestions but I chose not to add any PropertyKey to the MockInstance type. Somehow it does not feel right to say "this object has any key you can think of", e.g. it allows console.log(mock.foo) without any hint which is basically what we would get without typescript.

[mcous]

Thanks for the contribution! I couldn't find any existing ticket for this bug in Vitest - if one doesn't exist, it might be helpful to open one there, too, seeing as v3 is still the latest release