feat(expect): assertions (vitest-dev#278)

Co-authored-by: Anthony Fu <anthonyfu117@hotmail.com>

packages/vitest/src/index.ts

@@ -32,6 +32,7 @@ declare global {
       any(constructor: unknown): Any
       arrayContaining(expected: any): ArrayContaining
       stringMatching(expected: RegExp): StringMatching
+      assertions(expected: number): void
     }
 
     interface Assertion {

packages/vitest/src/integrations/chai/index.ts

@@ -1,5 +1,14 @@
 import chai from 'chai'
+import { getState, setState } from './jest-expect'
 
-export { assert, should, expect } from 'chai'
+export { assert, should } from 'chai'
 
-export { chai }
+const expect = (value: any, message?: string): Chai.Assertion => {
+  const { assertionCalls } = getState()
+  setState({ assertionCalls: assertionCalls + 1 })
+  return chai.expect(value, message)
+}
+
+Object.assign(expect, chai.expect)
+
+export { chai, expect }

packages/vitest/src/integrations/chai/jest-expect.ts

@@ -2,6 +2,35 @@ import type { Spy } from 'vitest'
 import { equals as asymmetricEquals, hasAsymmetric } from './jest-utils'
 import type { ChaiPlugin } from './types'
 
+type MatcherState = {
+  assertionCalls: number
+  expectedAssertionsNumber: number | null
+  expectedAssertionsNumberError: Error | null
+}
+const MATCHERS_OBJECT = Symbol.for('matchers-object')
+
+if (!Object.prototype.hasOwnProperty.call(global, MATCHERS_OBJECT)) {
+  const defaultState: Partial<MatcherState> = {
+    assertionCalls: 0,
+    expectedAssertionsNumber: null,
+    expectedAssertionsNumberError: null,
+  }
+  Object.defineProperty(global, MATCHERS_OBJECT, {
+    value: {
+      state: defaultState,
+    },
+  })
+}
+
+export const getState = <State extends MatcherState = MatcherState>(): State =>
+  (global as any)[MATCHERS_OBJECT].state
+
+export const setState = <State extends MatcherState = MatcherState>(
+  state: Partial<State>,
+): void => {
+  Object.assign((global as any)[MATCHERS_OBJECT].state, state)
+}
+
 // Jest Expect Compact
 // TODO: add more https://jestjs.io/docs/expect
 export const JestChaiExpect: ChaiPlugin = (chai, utils) => {
@@ -324,4 +353,19 @@ export const JestChaiExpect: ChaiPlugin = (chai, utils) => {
       callResult,
     )
   })
+
+  utils.addMethod(
+    chai.expect,
+    'assertions',
+    function assertions(expected: number) {
+      const error = new Error(`expected number of assertions to be ${expected}, but got ${getState().assertionCalls}`)
+      if (Error.captureStackTrace)
+        Error.captureStackTrace(error, assertions)
+
+      setState({
+        expectedAssertionsNumber: expected,
+        expectedAssertionsNumberError: error,
+      })
+    },
+  )
 }

packages/vitest/src/runtime/run.ts

@@ -4,6 +4,7 @@ import { vi } from '../integrations/vi'
 import type { ResolvedConfig, Suite, SuiteHooks, Task, Test } from '../types'
 import { getSnapshotClient } from '../integrations/snapshot/chai'
 import { hasFailed, hasTests, partitionSuiteChildren } from '../utils'
+import { getState, setState } from '../integrations/chai/jest-expect'
 import { getFn, getHooks } from './map'
 import { rpc, send } from './rpc'
 import { collectTests } from './collect'
@@ -41,7 +42,12 @@ export async function runTest(test: Test) {
 
   try {
     await callSuiteHook(test.suite, 'beforeEach', [test, test.suite])
+    setState({ assertionCalls: 0, expectedAssertionsNumber: null, expectedAssertionsNumberError: null })
     await getFn(test)()
+    const { assertionCalls, expectedAssertionsNumber, expectedAssertionsNumberError } = getState()
+    if (expectedAssertionsNumber !== null && assertionCalls !== expectedAssertionsNumber)
+      throw expectedAssertionsNumberError
+
     test.result.state = 'pass'
   }
   catch (e) {

test/core/test/jest-expect.test.ts

@@ -107,6 +107,20 @@ describe('jest-expect', () => {
     expect(complex).toHaveProperty('bar.arr.0')
     expect(complex).toHaveProperty('bar.arr.1.zoo', 'monkey')
   })
+
+  it('assertions', () => {
+    expect(1).toBe(1)
+    expect(1).toBe(1)
+    expect(1).toBe(1)
+    expect.assertions(3)
+  })
+
+  it('assertions with different order', () => {
+    expect.assertions(3)
+    expect(1).toBe(1)
+    expect(1).toBe(1)
+    expect(1).toBe(1)
+  })
 })
 
 it('timeout', () => new Promise(resolve => setTimeout(resolve, 500)))