疑似実装で理解するJest/Vitest spyOn

発端：なぜspyOn(obj.fn)ではなくspyOn(obj, 'fn')なのか

Jest/Vitestで関数の呼び出し履歴をテストしたいときにはspyOn関数が用いられる。オブジェクトobjに含まれる関数fnをテストする際には、vi.spyOn(obj, 'fn')という使い方となり、vi.spyOn(obj.fn)というインターフェースにはなっていない：

const spy = vi.spyOn(obj, 'fn') // vi.spyOn(obj.fn)じゃだめなのか？？

/* obj.fnを呼び出すことが期待される処理 */

expect(spy).toHaveBeenCalledWith('hoge')

この疑問を解消しつつspyOnがどのように動作しているかを理解するため、簡易版spyOnを実装してみた。

spyOnの仕組み

spyOnは大きく分けて以下2つのことをしている：

1. 渡された関数fnを、呼び出し履歴（など）を記憶するための層で包む
2. 包んだ関数で元の関数を"すり替える"

以下でそれぞれ詳しく解説する。

1. 呼び出し履歴を記憶する層で包む

元の関数を改変することなく呼び出し履歴を記憶するためには、元の関数を「スパイ機能」の層で包めば良い。例えば、呼び出したときの引数を逐次記録したい場合は

const spyOn = (fn) => {
  const mock = {
    calls: []
  }

  const spy = (...args) => {
    mock.calls.push(args)
    return fn(...args)
  }
  spy.mock = mock

  return spy
}

といった実装になる¹。元の関数の代わりにspyを呼び出せば、元の関数の呼び出しを"スパイ"することができる。実際、ここで作った簡易mock.callsはJest/Vitestと全く同じように動作する：

const spy = spyOn(obj.fn)

console.log(spy.mock.calls) // [] が出力される
spy('hoge', 1)
console.log(spy.mock.calls) // [['hoge', 1]] が出力される
spy('fuga', 2)
console.log(spy.mock.calls) // [['hoge', 1], ['fuga', 2]] が出力される

これを応用すればJest/Vitestにあるその他のインターフェースも実現できる。例えば実装をモックに置き換えるmockImplementationも

const spyOn = (fn) => {
  const spy = (...args) => {
    return fn(...args)
  }

  spy.mockImplementation = (newFn) => {
    fn = newFn
  }

  return spy
}

と簡単に実装できる。

余談：vi.mockとvi.spyOnの混乱ついて
mockImplementationの簡易実装で見たように、元々呼び出し回数をテストするために作ったはずのspyOnが、モックへの置き換えにも使えることがわかる。これは実用上便利なのだが、このせいでvi.mockとvi.spyOnのどちらを使ってもモックが実現できることになり、どちらを使えばいいのか混乱が生じているように思われる。片方しか使えないケースや記述量の差はあるものの、spyOnの裏側を理解するとどちらを使っても問題ない理由がわかると思う。

2. 包んだ関数で元の関数を"すり替える"

1.では呼び出し履歴の記録ができるようになっただけで、まだobj.fnは元の関数のままになっている。したがって、後は元の関数を1.で作ったスパイ済み関数にすり替えられればいい。イメージとしては

const spyOn = (fn) => {
  const mock = {
    calls: []
  }

  const spy = (...args) => {
    mock.calls.push(args)
    return fn(...args)
  }
  spy.mock = mock

  fn = spy // ★こんな感じですり替えたいが...
  return fn
}

spyOn(obj.fn)

としたいが、これはうまくいなかい。その理由は、spyOnの引数fnが参照している関数を置き換えたところで、obj.fnの参照先を置き換えたことにはならないからである。したがって、obj.fnをすり替えるspyOnのインターフェースは以下のようにならなくてはいけない：

const spyOn = (obj, functionName) => {
  const originalFunction = obj[functionName]
  const mock = {
    calls: []
  }

  const spy = (...args) => {
    mock.calls.push(args)
    return originalFunction(...args)
  }
  spy.mock = mock

  obj[functionName] = spy // プロパティの参照先を置き換える
  return spy
}

spyOn(obj, 'fn')

これがなぜspyOn(obj.fn)ではなくspyOn(obj, 'fn')としなければいけないかという問いに対する答えである。

まとめ

spyOnのインターフェースは一見奇妙であるものの、spyOnの疑似実装をしてみると必然的であることがわかった。spyOnの内部実装をイメージできると、テストがうまく動かないときの原因究明など役立つ場面があると思う。

1. mockを介してcallsを保持しているのはJest/Vitestの仕様と合わせるためなので、直接callsを保持しても本質的に違いはない。 ↩