Go单元测试 Mock 方案总结

2023-07-20 ⏳15.0分钟(6.0千字)

我从 2018 年转向 Go 语言开发,一晃就是五年。当年因为不了解 Go 生态,缺乏趁手的工具,所以选了一套很朴素但能用的方案做单元测试 Mock。后来伴随业务的发展,研发团队拆分成不同的小组。最近有的小组在做新项目,单元测试居然还在用五年前的临时方案。这确实让我意外。是时候总结一下过去的实践经验,做一拨分享了。

代码依赖

在我看来,单元测试的核心是 Mock。有了 Mock 才可以谈单元,不然所有的功能都耦合到一起,很难写测试用例。如果写用例很麻烦,长此以往下去,大家就不愿意、甚至压根不会去写测试用例。那么测试驱动开发自然就会烂尾,项目代码就会逐步失控。

我们先举一个具体的例子。比如说获取当前系统时间:

func Foo(t time.Time) {
  n := time.Now()
  if n.Sub(t) > 10 * time.Minute {
    // ...
  }
}

Foo() 函数接受一个参数t,会根据t与当前时间的差值的不同执行不同的业务逻辑。

我们很难给 Foo() 函数写测试用例,因为每次运行单元测试,它读到的当前时间都各不相同,结果也就不可预知,自然也就没办法通过程序来检测结果是否符合预期。这样的测试用例有的时候能跑过,有的时候就过不了,根本没法当成判断依据。

上面的示例虽然很简单,却非常典型。因为几乎所有的依赖都是这个样子,因为没办法控制底层函数的返回值,所以就无法检测上层函数的行为。类似的例子还包括:

除了上述这些明显的外部依赖,项目内部的依赖同样也很重要。现在的业务系统大多数都会分成几层。最简的的 MVC 也得有三层。如果采用了 DDD 模型,那可能会有四五层。我们的项目在参考了 DDD,但做了一定简化,共分成四层:

一个 RPC 函数可以操作多个 SVC 函数,一个 SVC 函数也可以操作多个 DAO 函数。从上到下单向依赖。有时 RPC 也可以直接操作 DAO 层函数。

这种分层结构导致了一个必然结果,那就是层级越低,依赖越少,功能相对越简单,写对应的单元测试也就越容易。随着层级的不断提高,功能和依赖会成指数方式增加。在写单测用例时如果不能有效隔离不同层之间的依赖,那么高层级代码就很难写对应的测试用例。几年实践下来发现,我们的测试用例大多都集中在 DAO 层,SVC 层就很少了,RPC 层几乎没有。究其原因,主要在于我们最初的方案不支持分层隔离依赖。

包装函数

那应该怎样 Mock 呢?让我们再回到上面的例子。

func Foo(t time.Time) {
  n := time.Now()
  if n.Sub(t) > 10 * time.Minute {
    // ...
  }
}

问题的根源是 Foo() 函数直接调用了 time.Now() 函数,而且我们无法控制 time.Now() 的返回值。Mock 的本质就是要控制 time.Now() 的行为。

最简单的办法就是自己写一个包装函数,比如:

now := time.Time{}

func myNow() time.Time {
  if !now.IsZero() {
    return now
  }
  return time.Now()
}

func Foo(t time.Time) {
  n := myNow()
  // ...
}

myNow() 函数封装 time.Now()myNow() 会优先检查 now 变量,如果非空则会返会 now,否则调用 time.Now() 返回当前实际的时间。我们写测试用例的时候可以通过设置 now 的值来调整系统时间,从而控制 Foo() 的行为。

大家可能会觉得这种方案太 Low 了,不但需要封装已有的函数,还要提供包一级的全局变量,这是妥妥的反模式呀!但是这种处理其实还挺常见的。Go 语言就是这样,不管白猫黑猫,能抓老鼠才是王道。

典型的就是 net/http 中的 DefaultTransport 变量,我们可以设置自己的实现,从而影响 http.Get() 等标准库函数的行为。还有一个专门 Mock HTTP 请求的包,叫 jarcoal/httpmock1,底层就是通过修改 http.DefaultTransport 实现 Mock 功能。

这种方案除了丑之外,还有一个缺点,不能并发运行。如果有多个测试用例并发修改全局变量,则可能产生非预知结果。所以就有了第二种方案,依赖注入2

依赖注入

简单来说就是不直接调用依赖的函数,而是把依赖当成普通变量传入上层函数。我们可以把上面的代码做如下改写:

type Nower func() time.Time

func Foo(t Time.Time, now Nower) {
  n := now()
  // ...
}

// 调用传参
Foo(t, time.Now)

这里把 time.Now 当成参数传给 Foo() 函数。如果想自定义系统时间,可以传入新的函数:

Foo(t, func() time.Time {
  s := "2021-05-20T15:34:20Z"
  t, _ := time.Parse(time.RFC3339, s)
  return t
})

为了单元测试而单独传入 Nower 参数,怎么看怎么捌扭。但在很多实际场景下,这种方案就显得很自然,甚至是优雅。比如下例:

func ReadContents(f io.ReadCloser, numBytes int) ([]byte, error) {
  defer f.Close()
  data := make([]byte, numBytes)
  return f.Read(data)
}

func Foo() {
  f, _ := os.Open("foo.txt")
  data, err := ReadContents(f, 50)
  // ...
}

这里 ReadContents() 函数第一个入参是 io.ReadCloser 接口,业务代码 Foo() 传入的是 *os.File 指针。我们如果想单独测试 ReadContents() 函数,则可以自己实现特殊的 io.ReadCloser 并传给 ReadContents() 函数。

type MockFile struct {}

// 模拟读取报错情形
func (f *MockFile) Read(p []byte) (n int, err error) {
  return 0, errors.New("foo")
}
func (f *MockFile) Close() error { return nil }

b, err := ReadContents(&MockFile{}, 10)
// ...

这里的优雅更多源自接口抽象。这里的自然更多是因为 f 就是 ReadContents() 要用的参数,不是前面那种为了测试而专门添加的额外参数。

但这种方案也不是没有缺点。第一个缺点就是写类似 MockFile 之类的实现并不容易。尤其是当接口有很多函数的时候,为了 Mock 某几个函数,我们要把接口中的所有函数都实现一遍。

解决这个问题也有两个办法。第一个是在实现中直接嵌入对应的接口,然后只实现需要的函数。比如:

type Foo interface {
	A(i int) int
	B(i int) int
}

type Bar struct {
	Foo // 直接嵌入接口
}

func (b *Bar) A(i int) int { return i+1 }

var f Foo = &Bar{}

func Test(f Foo) {
  f.A(1) // output 2
  f.B(1) // panic
}

这里接口 FooA()B() 两个函数。结构体 Bar 只实现了 A() 函数,但因为直接嵌入了接口 Foo,所以也被当成了 Foo 的实现。不过这取巧的办法也不是万能的。如果调用没有实现的函数会直接 panic 退出。

如果嫌自己实现太麻烦,也可以使用一些自动生成代码的方案。当前最有名的要属 Go 语言官方提供的 golang/mock3 了。但很可惜,从 2023 年 6 月起,该项目已经烂尾。但它的思路还值得学习,而且 Uber 还在维护一个 fork 版本。

先安装 mockgen 工具:

go install github.com/golang/mock/mockgen@latest

然后生成 mock 代码:

mockgen -source=foo.go -package=foo > mock_gen.go

它会为上面的 Foo 接口生成 NewMockFoo() 函数。在单元测试中可以这样使用:

func TestFoo(t *testing.T) {
  ctrl := gomock.NewController(t)
  defer ctrl.Finish()

  m := NewMockFoo(ctrl)

  m.
    EXPECT().
    A(gomock.Eq(1)). // 断言入参为一
    Return(2) // 返回二

  Test(m)
}

gomock 生成的实现有很多特性,最常用的是检查函数的入参和设置反回值。其他功能请参考官方文档,在此就不展开了。但本质上是为接口添加特殊的实现,跟前面手写代码没有区别。

依赖注入的另一个缺点是侵入性很强,需要为了注入依赖而改写业务代码。当依赖非常负杂的时个,这种改写就很难维护,所以又出现了各种依赖注入框架。比较知名的是 Google 开源的 wire 框架4

为了将就注入框架,很多业务框架也做了不同程度的妥协。如果套用我们前面说的业务分层模型,有些框架为了利用 wire 框架,它们会在每一层定义一个或者多个很大的接口,然后分别注册到 wire 框架,并且让 wire 管理依赖关系。这种处理方式除了难以理解外,另一个突出的缺点就是让业务代码倾向于写很大的接口。我曾经见过有业务项目把所有的 DAO 层函数都放到一个接口。有些新框架采用 DDD 模式,但很多 Repository 也是非常大。

每个 DAO 或者 Repo 都要声明接口,每个接口只有一个实现,而且都是全局单例。所有接口和实现都通过 wire 管理。所有的这一切都为一个目的就是依赖注入。而依赖注入的目标就是方便写测试用例。就我而言,这种为了写单元测试而如此大费周折实在有点得不偿失。所以我一直拒绝这一类方案。

猴子补丁

我想要的一定是自然的方案。虽然我们要实现测试驱动开发,但这只能是从逻辑角度来看。从实现角度,我们不能为了方便写单元测试而大幅调整业务代码的自然顺序。前文提到的依赖注入也叫控制反转。反转了就不自然。

怎么才算自然呢?回到最开始的例子:

func Foo(t time.Time) {
  n := time.Now()
  if n.Sub(t) > 10 * time.Minute {
    // ...
  }
}

业务代码就是需要直接调用 time.Now(),不要整任何幺蛾子,我不想为了单元测试修改业务代码。但我还要在测试的过程中调整 time.Now() 的返回值。这就是既要、又要、还要。那能实现吗?当然能,但需要一点黑科技,这就是 Monkey Patch 技术,也叫猴子补丁。

市面上这一类框架有很多变种。但核心思想都源自 Bouke 大神,他首次提出在 Go 语言中通过动态修改函数代码段来实现 Mock 功能。我也写过多篇介绍文章5介绍其实现原理。本文只讲使用方法。

就上面的例子来说,我们要 Mock time.Now() 函数,可以直接做如下操作:

import "github.com/go-kiss/monkey"

func TestFoo(t *testing.T) {
  // 直接修改 time.Now 的返回值
  defer monkey.Patch(time.Now, func() time.Time {
    s := "2021-05-20T15:34:20Z"
    t, _ := time.Parse(time.RFC3339, s)
    return t
  }).Unpatch()

  Foo(t1)
}

从效果上来看,这里通过 monkey.Patch() 函数可以直接修改 time.Now() 的返回值。但实际上是框架在运行时改掉了 time.Now() 对应的机器码,让函数跳转到了我们传入的匿名函数。

这个方案就非常优雅了。不用写包装函数,不用注入依赖,不用生成代码。什么也不用做。业务代码该怎么写就怎么写。无论被调的是自己写的函数还是标准库或者三方库提供的函数,都可以使用 Monkey Patch 来修改它们的行为。这才是理想的 Mock 工具该有的样子,简单易用零侵入!

但是天下没有免费的午餐。Monkey Patch 这么好,不可能没有缺点吧。当然有,而且缺点还很严重,要不怎么说是黑科技呢?

首先,在程序运行时修改函数的代码段这个操作有些操作系统本身就不支持。其中就有苹果的 M1 芯片平台。苹果是真的贱,为了所谓的安全性,在系统层面禁止内存同时拥有可写和可执行权限。所以 Monkey Patch 没法在 M1 设备上原生运行。这是硬伤。但好在 M1 可以模拟 x86 架构,运行时指定 GOARCH=AMD64 就可以临时绕过这个问题。

其次就是兼容性问题。因为底层实现依赖 Go 语言的 ABI,但 Go 语言只保证 API 下向兼容,从不保证 ABI 稳定不变。会不会未来某个 Go 版本 ABI 大改,导致之前写的单元测试都没法正常运行了呢?有这个可能,这也是我当年最大的顾虑。但经过研究发现,ABI 中只依赖一个额外的寄存器。Go ABI 中正好保留了 R12 和 R13 用作通用场景,语言本身不会直接使用。所以我们用它们保存临时数据不会产生问题。而且 Go 在 1.16 版本开始改用寄存器传递函数参数6,这样的 ABI 改动也没有影响到 Monkey Patch 方案,所以也就没有什么好顾虑的了。

第三个问题是不支持 Mock 接口。如果我们只持有 io.Reader 变量,没办法 Mock 它的 Read() 函数,只有找到底层的实现对象才可以。

第四个问题是可能有并发安全问题。从原理上讲,Monkey Patch 的实质是修改某函数指针指向的内存内容(内容为机器码)。因为函数指针全局唯一,如果有多个协程同时修改就有可能产生并发问题。但考虑到只在运行单元测试时使用,所以问题不大。而且几年实践下来也没碰到过这类的问题。

最后一个问题就是跟最前面的包装函数方案一样,也不支持并发运行测试用例。这本质上也是因为函数指针全局唯一。一个协程改了某函数的机器码,另一个协程可能就会受到影响。为了解决这个问题,我花了很大的力气研究 Monkey Patch 的原始实现,最终设计了一套支持协程隔离的增强版本 go-kiss/monkey7。基本上解决了这个问题。

前面的例子其实用的就是我的改造版本。go-kiss/monkey 默认开启协程隔离,不同协程可以分别 Mock 同一个函数而互不影响。在有些特殊场景下需要全局 Mock,可以传入专门的参数:

monkey.Patch((*net.Dialer).DialContext, func(_ *net.Dialer, _ context.Context, _, _ string) (net.Conn, error) {
  return nil, fmt.Errorf("no dialing allowed")
}, monkey.OptGlobal) // global 参数

_, err := http.Get("http://taoshu.in")
fmt.Println(err) // Get http://taoshu.in: no dialing allowed

但我们能说 Monkey Patch 就是终极方案吗?显然也不能。Go 语言的 ABI 变化始终像一把利剑悬在头上,说不定哪天直有可能斩杀所有的测试用例。而且它也不能原生支持苹果的 M1 芯片。难道就没有两全其美的办法了吗?

代码改写

完美的没找到,但找到了一支潜力股,那就是xhd2015/go-mock项目8。go-mock 受 go tool cover 的启发9实现了一套在编译期改写被 Mock 函数的功能。

假如我们有一个函数 Add()供其他业务代码调用:

func Add(a, b int) int {
  return a + b
}

那么 go-mock 在编译前会将其改写成:

package hello;import _mock "github.com/xhd2015/go-mock/mock"

import "fmt"

func Add(a,b int) {var _mockreq = struct{}{};var _mockresp struct{};_mock.TrapFunc(nil,&_mock.StubInfo{PkgName:"github.com/xhd2015/go-mock/tmp",Owner:"",OwnerPtr:false,Name:"Add"}, nil, &_mockreq, &_mockresp,_mockHello,false,false,false);}; func _mockAdd(a,b int) int {
  return a + b
}

这里的核心是在原来的 func Add(a,b int){ 插入了一堆神奇的代码,换行后如下:

func Add(a,b int) {
  var _mockreq = struct{}{}
  var _mockresp struct{}
  _mock.TrapFunc(nil,&_mock.StubInfo{PkgName:"github.com/xhd2015/go-mock/tmp",Owner:"",OwnerPtr:false,Name:"Add"}, nil, &_mockreq, &_mockresp,_mockHello,false,false,false)
}
func _mockAdd(a,b int) int {
  return a + b
}

其实是把原来的 Add() 折成了两部分,原来的函数逻辑变成 _mockAdd();原来的函数入口被改写成调用 _mock.TrapFunc() 函数,这个函数会根据单元测试的 Mock 规则调用对应的 Mock 函数。把所有的代码合成一行放插到原函数后面是为了不影响单元测试报错时的行号或者调用栈。该技巧可以参考上面提到的 Go 语言官方博客。go-mock 最核心的命令是 rewrite,实现比较复杂,但思路很明确,就是扫描当前项目用到的所有包,从中找到所有函数,然后在后面插入对应的 Mock 代码。

现在给大家一个完整的示例,官方给的示例不容易理解。

先安装 go-mock 工具

go install github.com/xhd2015/go-mock

假设我们的项目结构如下:

├── bar
│   ├── bar.go
│   └── bar_test.go
├── foo
│   └── foo.go
├── go.mod
├── go.sum
└── main.go

foo.go 中定义如下函数:

package foo

import (
  "context"
)

func Hi(ctx context.Context, a, b int) int {
  return a + b
}

bar.go 中调用 foo.Hi() 函数:

package bar

import (
  "context"
  "hello/foo"
)

func Hi(ctx context.Context, a, b int) int {
  return foo.Hi(ctx, a, b) + 1
}

现在我们写 bar.Hi() 函数的测试用例。但在写用例之前我们需要生成测试专用的代码:

go-mock rewrite -f -v ./...

它默认会在 test 目录为每一个源文件生成对应的 Mock 辅助代码文件:

test
└── mock_gen
    ├── bar
    │   └── mock.go
    ├── foo
    │   └── mock.go
    └── mock.go

这里面的代码我们在写测试用例的时候会用到。同时,它还会在一个临时目录(打开-v 会输出到终端)生成前面说的改写后的源码文件,也就是插入 TrapFunc 的文件。这部分我们可以不去管它。现在我们开始写测试用例。

写之前还要导入 go-mock 包:

go get github.com/xhd2015/go-mock

然后添加测试代码 bar_test.go

package bar

import (
  "context"
  "fmt"
  "testing"
  
  // 导入生成的辅助包
  mock_foo "hello/test/mock_gen/foo"
)

func TestHi(t *testing.T) {
  ctx := context.Background()
  // 设置 Mock 函数
  ctx = mock_foo.Setup(ctx, func(m *mock_foo.M) {
    // 每个函数都有对应的变量
    m.Hi = func(ctx context.Context, a, b int) int {
      return a - b
    }
  })
  // 此时调用 Hi() 底层会调用刚才 Mock 过的函数
  fmt.Println(Hi(ctx, 1, 2))
}

不能直接使用 go test 运行,而要改用 go-mock:

go-mock test hello/bar

会看到输出结果为 12+1=01-2+1=0,而不是原来的 1+2+1=41+2+1=4。Mock 成功。

现在总结一下 go-mock 的优点。

首先,它是纯 Go 语言实现,几乎没有黑科技,对平台没有任何要求,也不依赖 Go 语言的 ABI。所以它的兼容性最好,支持所有平台,也几乎支持所有 Go 语言版本。因为是源码级依赖,后续新版本的 Go 语言也几乎不会对 go-mock 产生影响。这是它最大的亮点!

其次,它没有并发安全问题,因为所有 Mock 状态都是通过 ctx 来传递,这也是官方推荐的标准做法。

基于以上两点,我判断 go-mock 是一支潜力股,未来大有可为。而且我也希望后续能参与到这个项目。

但 go-mock 就完美无瑕了吗?当然不是!而且缺点也同样明显。这些缺点也只是我通过有限的测试和研究发现的,也可能是我不了解,或者后续能解决。这里先列出来供大家参考。

第一,它只能 Mock 自己写的代码。像 time.Now() 这种函数就没办法 Mock。就我目前的了解,它好像也没办法 Mock 第三方包中的函数。

第二,要 Mock 的函数必须传入 ctx 参数。这个问题不是很大,因为现在 Go 生态中几乎大部分函数都支持传入 ctx 了。但仍然有不少场景不传这个参数。不传就没法 Mock,确实有点小难受。也正是因为依赖 ctx,所以它目前不支持 Mock 那些会返回 ctx 的函数。这也算是并发安全的小代价吧。

第三,扫描并生成代码可能会拖慢测试速速度。

第四,不支持 Mock 泛型函数。这个同样问题不大。

第五,当前实现不支持 Mock 当前包内的函数,会产生循环依赖。但应该能解决。

总得来说 go-mock 前景非常不错。因为业务项目大概率不会直接调用标准函数或者第三方包,而是会提供一层封装。这样的话正好绕开了 go-mock 的缺点。项目不同层级的代码可以随意 Mock,已经可以解决大部分问题了。用 2-8 原则来看,go-mock 确实是非常不错的工具。

我个人的建议是两种方案都可以尝试。Monkey Patch 有一定风险,但更灵活;go-mock 一点风险也没有,但 Mock 场景有一定的限制。两者暂时无法互为替换,只有靠大家根据自己的业务场景和团队现实来选择了。


写到这本文的主要内容就该结束了。但文中开头提到早期方案居然没有合适的位置摆放。为了内容的完整性,我把它放到最后,作为全文的补充供大家参考。

早期方案

我在前面说过,因为在 2018 年我对 Go 语言本身都缺乏足够的了解,更别提 Monkey Patch 这种黑科技方案了。而且那时候也无力实现协程隔离,如果用的话也就无法开启并行测试。

在当时的条件下,我们选择了一种土味实足的方案——功能模拟。

所谓功能模拟就是缺什么补什么。代码需要数据库我们就跑一个 MySQL 实例;代码需要缓存我们就跑一个 Redis 实例;代码需要调用外部 API,这就没办法了,我们用 jarcoal/httpmock 来 Mock。为了让不同的 Pipeline 互不干扰,我们采用 Docker 模式运行 Pipeline。为了让单元测试环境跟代码保持一致,我们在运行用例之前会从开发数据库自动导入表结构。为了生成测试数据,我们通过专门的 seed.sql 往数据库是灌数据。

这一套系统只实现了一个目标,那就是不同的 Merge Request 的单元测试可以并行运行,互不干扰。但同一个分支的 Pipeline 因为共用一套数据持久层,所以只能串行执行。五年下来导致我们的测试用例越来越慢。另一个明显的问题是大多数测试用例都在 DAO 层,因为这一层离数据最近,编写测试用例也最容易。越上层的代码依赖越多,用例也越难写,自然也就越来越少。

团队拆分之后,当其他组继续做新服务的时候,大家居然还沿用了这套方案。要不是因为中间用 IT 环境调整出现无法导出表结构的问题,我都不知道这件事。我之前跟大家分享的 Mock 方案几乎没人理会,业界新的方案更是无人问津。当老方案比较慢的时候,大家首先想到的就是在老方案上做改进。一种改进的方案是给不同的 package 标记不同的表结构依赖,然后在启动前为不同的 package 创建不同的数据库,以此来实现数据隔离,进而开启并发测试。

这种方案不是不行,但总给人一种螺蛳壳里做道场的感觉。有感而发,故成此文。


  1. https://github.com/jarcoal/httpmock↩︎

  2. 依赖注入不只解决并发运行单测的问题,但本文略过其他方面↩︎

  3. https://github.com/golang/mock↩︎

  4. https://github.com/google/wire↩︎

  5. ./monkey/monkey.html↩︎

  6. https://github.com/golang/go/issues/40724↩︎

  7. 具体原理可以参考我的专门文章 ./monkey/monkey-2.html↩︎

  8. https://github.com/xhd2015/go-mock↩︎

  9. https://go.dev/blog/cover↩︎