Skip to content

Commit

Permalink
feat(spy): implement mock.withImplementation API (#2835)
Browse files Browse the repository at this point in the history
Co-authored-by: Vladimir Sheremet <sleuths.slews0s@icloud.com>
  • Loading branch information
obadakhalili and sheremet-va authored Feb 24, 2023
1 parent 6ff6c6e commit 610b1d4
Show file tree
Hide file tree
Showing 4 changed files with 264 additions and 81 deletions.
37 changes: 37 additions & 0 deletions docs/api/mock.md
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,43 @@ You should use spy assertions (e.g., [`toHaveBeenCalled`](/api/expect#tohavebeen
console.log(myMockFn(), myMockFn(), myMockFn(), myMockFn())
```

## withImplementation

- **Type:** `(fn: Function, callback: () => void) => MockInstance`
- **Type:** `(fn: Function, callback: () => Promise<unknown>) => Promise<MockInstance>`

Overrides the original mock implementation temporarily while the callback is being executed.

```js
const myMockFn = vi.fn(() => 'original')
myMockFn.withImplementation(() => 'temp', () => {
myMockFn() // 'temp'
})
myMockFn() // 'original'
```

Can be used with an asynchronous callback. The method has to be awaited to use the original implementation afterward.

```ts
test('async callback', () => {
const myMockFn = vi.fn(() => 'original')
// We await this call since the callback is async
await myMockFn.withImplementation(
() => 'temp',
async () => {
myMockFn() // 'temp'
},
)
myMockFn() // 'original'
})
```

Also, it takes precedence over the [`mockImplementationOnce`](https://vitest.dev/api/mock.html#mockimplementationonce).

## mockRejectedValue

- **Type:** `(value: any) => MockInstance`
Expand Down
33 changes: 32 additions & 1 deletion packages/spy/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ export interface SpyInstance<TArgs extends any[] = any[], TReturns = any> {
getMockImplementation(): ((...args: TArgs) => TReturns) | undefined
mockImplementation(fn: ((...args: TArgs) => TReturns) | (() => Promise<TReturns>)): this
mockImplementationOnce(fn: ((...args: TArgs) => TReturns) | (() => Promise<TReturns>)): this
withImplementation<T>(fn: ((...args: TArgs) => TReturns), cb: () => T): T extends Promise<unknown> ? Promise<this> : this
mockReturnThis(): this
mockReturnValue(obj: TReturns): this
mockReturnValueOnce(obj: TReturns): this
Expand Down Expand Up @@ -208,6 +209,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
}

let onceImplementations: ((...args: TArgs) => TReturns)[] = []
let implementationChangedTemporarily = false

let name: string = (stub as any).name

Expand Down Expand Up @@ -248,6 +250,35 @@ function enhanceSpy<TArgs extends any[], TReturns>(
return stub
}

function withImplementation(fn: (...args: TArgs) => TReturns, cb: () => void): EnhancedSpy<TArgs, TReturns>
function withImplementation(fn: (...args: TArgs) => TReturns, cb: () => Promise<void>): Promise<EnhancedSpy<TArgs, TReturns>>
function withImplementation(fn: (...args: TArgs) => TReturns, cb: () => void | Promise<void>): EnhancedSpy<TArgs, TReturns> | Promise<EnhancedSpy<TArgs, TReturns>> {
const originalImplementation = implementation

implementation = fn
implementationChangedTemporarily = true

const reset = () => {
implementation = originalImplementation
implementationChangedTemporarily = false
}

const result = cb()

if (result instanceof Promise) {
return result.then(() => {
reset()
return stub
})
}

reset()

return stub
}

stub.withImplementation = withImplementation

stub.mockReturnThis = () =>
stub.mockImplementation(function (this: TReturns) {
return this
Expand Down Expand Up @@ -275,7 +306,7 @@ function enhanceSpy<TArgs extends any[], TReturns>(
stub.willCall(function (this: unknown, ...args) {
instances.push(this)
invocations.push(++callOrder)
const impl = onceImplementations.shift() || implementation || stub.getOriginal() || (() => {})
const impl = implementationChangedTemporarily ? implementation! : onceImplementations.shift() || implementation || stub.getOriginal() || (() => {})
return impl.apply(this, args)
})

Expand Down
Loading

0 comments on commit 610b1d4

Please sign in to comment.