Mock in Go Unit Testing

I’ve been developing in Go since 2018, and it’s been five years since then. I didn’t understand the Go ecosystem and lacked tools, so I chose a very simple but usable solution to do unit testing Mock, and then along with the development of the business, the R&D team was split into different groups. Recently, some teams are working on new projects, but the unit tests are still using the same temporary solution from five years ago. This really surprised me. It’s time to summarise the past experience and share it.

Code Dependencies

In my opinion, the core of unit testing is Mock, with Mock you can talk about unit, otherwise all the functions are coupled together, it is difficult to write test cases. If writing test cases is troublesome, in the long run, people will not be willing to write test cases, or even will not go to write test cases at all. Then test-driven development will naturally suck, and the project code will gradually get out of control.

Let’s take a specific example. For example, get the current system time:

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

The Foo() function takes one argument, t, and performs different business logic depending on the difference between t and the current time.

It’s hard to write test cases for the Foo() function because every time you run a unit test, it reads a different current time, the result is unpredictable, and there’s no way to programmatically check that the result is as expected. Such test cases can sometimes pass, sometimes not, and can’t be used as a basis for judgement.

The above example is simple but very typical. Because almost all dependencies look like this, there is no way to detect the behaviour of the upper level function because there is no way to control the return value of the underlying function. Similar examples include:

calling external HTTP interfaces via net/http
Querying a relational database via database/sql.
Reading and writing the Redis cache via redis/go-redis.
……

Apart from these obvious external dependencies mentioned above, the internal dependencies of the project are equally important. Most business systems nowadays are divided into several layers. The simplest MVC has three layers. If the DDD model is used, there may be four or five layers. In our project, we referenced DDD, but simplified it a bit and divided it into four layers:

RPC interface layer, including interface description IDL, preliminary parameter conversion and checking logic, etc.
SVC service layer, provides a complete implementation of business scenarios for the RPC layer, relying on multiple DAO layer components.
DAO data access layer, encapsulate the details of the underlying data acquisition, and provide upward access to the SVC layer.
PKG tooling layer, providing infrastructure or tooling components.

An RPC function can operate on multiple SVC functions, and an SVC function can operate on multiple DAO functions. Unidirectional dependency from top to bottom. Sometimes RPC can also operate directly on DAO layer functions.

This hierarchical structure leads to a corollary, that is, the lower the level, the fewer the dependencies, the relatively simpler the function, write the corresponding unit tests will be easier. As the tiers increase, the functionality and dependencies increase exponentially. If you can’t effectively isolate dependencies between different layers when writing unit tests, it will be difficult to write corresponding test cases for higher level code. After several years of practice, we found that most of our test cases are concentrated in the DAO layer, very few in the SVC layer, and almost none in the RPC layer. The reason for this is that our initial solution does not support isolation of dependencies by layers.

So how should we Mock it? Let’s go back to the example above.

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

The root of the problem is that the Foo() function calls the time.Now() function directly, and we have no control over the return value of time.Now(). The essence of Mock is to control the behaviour of time.Now().

The easiest way to do this is to write your own wrapper function, for example:

now := time.Time{}

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

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

Wrap time.Now() with the myNow() function. myNow() checks the now variable first and returns now if it is not empty, otherwise it calls time.Now() to return the actual current time. We can control the behaviour of Foo() by setting the value of now to adjust the system time when writing test cases.

You may think that this solution is too low, not only do you need to encapsulate existing functions, but you also need to provide global variables at the package level, which is a proper anti-pattern! But this is actually quite common.

Typical is the DefaultTransport variable in net/http, where we can set our own implementation to affect the behaviour of standard library functions such as http.Get(). There’s also a package for mocking HTTP requests called jarcoal/httpmock, which underneath implements the mock functionality by modifying http.DefaultTransport.

This solution, besides being ugly, has the disadvantage that it cannot be run concurrently. If there are multiple test cases concurrently modifying global variables, unpredictable results may occur. So comes the second option, dependency injection.

Dependency Injection

This simply means that instead of calling the dependent function directly, the dependency is passed into the higher level function as a normal variable. We can rewrite the above code as follows:

type Nower func() time.Time

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

// call
Foo(t, time.Now)

Here time.Now is passed as an argument to the Foo() function. If you want to customise the system time, you can pass in a new function:

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

Passing in a separate Nower parameter for the purpose of unit testing is a twisted way to look at it. But in many real-world scenarios, this solution seems natural, even elegant. Take the following example:

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)
  // ...
}

Here the ReadContents() function’s first input parameter is the io.ReadCloser interface, and the business code Foo() passes in the *os.File pointer. If we want to test the ReadContents() function separately, we can implement our own special io.ReadCloser and pass it to the ReadContents() function.

type MockFile struct {}

// Simulation of read exception
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)
// ...

The elegance here comes more from the interface abstraction. The naturalness here comes more from the fact that f is the parameter to be used by ReadContents(), not the extra parameter that was added specifically for testing purposes.

But this solution is not without its drawbacks. The first drawback is that it’s not easy to write implementations like MockFile. Especially if the interface has a lot of functions, we have to implement all the functions in the interface in order to Mock some of them.

There are two ways to solve this problem. The first is to embed the corresponding interface directly in the implementation, and then implement only the functions you need. For example:

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

type Bar struct {
    Foo // embedding
}

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
}

Here the interface Foo has two functions A() and B(). The structure Bar only implements the A() function, but since it’s directly embedded in the interface Foo, it’s also treated as an implementation of Foo. This approach is not foolproof, though. If you call a function that doesn’t implement A(), it will be panic.

If it’s too much of a hassle to do it yourself, you can also use some automated code generation schemes. The most famous one is golang/mock, which is officially available for the Go language. Unfortunately, as of June 2023, this project is dead. But the idea is still worth learning, and Uber still maintains a fork version.

Start by installing the mockgen utility:

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

Then generate the mock code:

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

It generates the NewMockFoo() function for the Foo interface above. You can use it like this in unit tests:

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)
}

The implementation that gomock generates has a number of features, the most common being checking incoming arguments to functions and setting back values. Please refer to the official documentation for other features, so I won’t expand on them here. However, it is essentially adding a special implementation to an interface, which is no different from the handwritten code in the previous section.

Another disadvantage of dependency injection is that it is very intrusive and you need to rewrite the business code in order to inject the dependency. When the dependency is very negative miscellaneous times, this rewrite is very difficult to maintain, so there is a variety of dependency injection framework. The more well-known is Google’s open source wire framework.

Many business frameworks also make compromises of varying degrees in order to accommodate the injection framework. If we apply the business layering model we described earlier, some frameworks will define one or more very large interfaces at each layer in order to take advantage of the wire framework, then register them with the wire framework and let the wire manage the dependencies. In addition to being difficult to understand, another prominent drawback of this approach is that business code tends to write very large interfaces. I’ve seen business projects put all the DAO layer functions into one interface. Some of the newer frameworks use the DDD model, but many of the Repositories are also very large.

Each DAO or Repo has to declare interfaces, each interface has only one implementation, and they are all global singletons. All interfaces and implementations are managed through wires. All of this is done for the purpose of dependency injection. The goal of dependency injection is to make it easier to write test cases. As far as I’m concerned, this kind of effort to write unit tests just doesn’t pay off. So I’ve always rejected this type of solution.

Monkey Patch

What I want must be the natural solution. While we want to implement test-driven development, this can only be done from a logical perspective. From an implementation point of view, we can’t drastically adjust the natural order of business code to make it easier to write unit tests. The previously mentioned dependency injection is also called control inversion. Inversion is not natural.

How is it natural? Go back to the very first example:

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

The business code just needs to call time.Now() directly without any fuss, and I don’t want to change the business code for the unit tests. But I also want to tweak the return value of time.Now() during testing. Can that be achieved? Sure, but it takes a little bit of hacking, and that’s where Monkey Patch comes in.

There are many variants of this type of framework on the market. But the core idea comes from Bouke, who first proposed to implement Mock functionality in Go by dynamically modifying function snippets. I have also written a number of articles to introduce the principle of its implementation. In this article, we will only talk about the usage.

In the above example, we want to Mock time.Now() function, we can directly do the following:

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

func TestFoo(t *testing.T) {
  // Modify the return value of time.Now directly
  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)
}

Effectively, the monkey.Patch() function can be used to modify the return value of time.Now() directly. But in reality, the framework is changing the machine code of time.Now() at runtime, so that the function jumps to the anonymous function we passed in.

This is a very elegant solution. No wrapper functions to write, no dependencies to inject, no code to generate. You don’t have to do anything. The business code is written as it should be. Whether the function being called is one you wrote yourself or one provided by a standard or third-party library, you can use Monkey Patch to modify their behaviour. This is the ideal Mock tool should look like, easy to use zero intrusion!

But there’s no such thing as a free lunch, and if Monkey Patch is so good, it can’t be without its drawbacks. Of course, there are, and the drawbacks are very serious, otherwise it would be called a black technology.

First of all, the operation of modifying the code segment of a function at runtime is not supported by some operating systems. Among them is Apple’s M1 chip platform. Apple is really cheap, for the sake of the so-called security, in the system level to prohibit the memory at the same time can have write and executable permissions. So Monkey Patch can’t run natively on M1 devices. That’s the hard part. But the good news is that the M1 can emulate the x86 architecture, and you can temporarily get around this by specifying GOARCH=AMD64 at runtime.

The second issue is compatibility. Because the underlying implementation relies on the Go language ABI, Go only guarantees downward compatibility of the API, but never guarantees that the ABI is stable and unchanged. Will the ABI change in some Go version in the future, so that all the unit tests written before won’t work properly? This is a possibility, and it was my biggest concern back then. But after some research, I found that the ABI only relies on one extra register, and the Go ABI retains R12 and R13 for general-purpose scenarios, which are not directly used by the language itself. So we can use them to save temporary data without any problems. Moreover, Go has changed to use registers to pass function arguments since version 1.16, and this ABI change does not affect the Monkey Patch scenario, so there is nothing to worry about.

The third issue is that the Mock interface is not supported. If we only hold the io.Reader variable, there’s no way to mock its Read() function, only find the underlying implementation.

The fourth problem is that there may be concurrency safety issues. In principle, a Monkey Patch essentially modifies the contents of the memory pointed to by a function pointer (the contents are machine code). Since function pointers are globally unique, if more than one goroutine modifies them at the same time, there may be concurrency issues. But considering that it’s only used when running unit tests, it’s not a big problem. And I haven’t encountered this kind of problem in several years of practice.

The last problem is that, like the wrapped function solution at the top, it doesn’t support running test cases concurrently. This is essentially because function pointers are globally unique. If a goroutine changes the machine code of a function, another worker may be affected. To solve this problem, I put a lot of effort into researching the original implementation of Monkey Patch, and ended up designing a set of enhancements to go-kiss/monkey that support concurrent isolation. This basically solves the problem.

The previous example actually uses my modified version. By default, go-kiss/monkey enables goroutine isolation, so that different goroutines can mock the same function without affecting each other. In some special cases where you need a global mock, you can pass in a special parameter:

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

But can we say that Monkey Patch is the ultimate solution? No. The Go language’s ABI changes are always a sword hanging over your head, threatening to kill all your test cases one day. And it doesn’t natively support Apple’s M1 chip. Isn’t there a way to get the best of both worlds?

Code rewrite

We didn’t find the perfect one, but we did find a potential one, which is the xhd2015/go-mock project;go-mock is inspired by the go tool cover, and implements a set of features to rewrite mocked functions at compile time.

Suppose we have a function Add() that is called by other business code:

Then go-mock rewrites it before compiling:

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
}

The core is a bunch of magical code inserted into the original func Add(a,b int) and {, with line breaks as follows:

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
}

In fact, the original Add() is folded into two parts, the original function logic becomes _mockAdd(); the original function entry is rewritten to call _mock.TrapFunc() function, this function will call the corresponding Mock function according to the Mock rules of the unit test. The reason for putting all the code into one line and inserting it after the original function is so that it doesn’t affect the line number or the call stack in case of a unit test error, you can refer to the Go language official blog mentioned above. go-mock’s core command is rewrite, which is a bit more complicated to implement, but the idea is very clear, that is, to scan all packages used in the current project, find all the functions from them, and then insert the corresponding Mock code after them. and then insert the corresponding Mock code after it.

Now I’m going to give you a complete example, the official example is not easy to understand.

Start by installing the go-mock utility:

`1`	`go install github.com/xhd2015/go-mock`

Suppose our project is structured as follows:

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

The following function is defined in foo.go:

package foo

import (
  "context"
)

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

Call the foo.Hi() function in bar.go:

package bar

import (
  "context"
  "hello/foo"
)

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

Now we write the test case for the bar.Hi() function. But before we write the test case we need to generate test specific code:

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

By default, it generates a Mock helper code file in the test directory for each source file:

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

This is the code we will use when writing test cases. It also generates the rewritten source code file mentioned earlier in a temporary directory (opening -v will output it to the terminal), which is where TrapFunc is inserted. We can leave this part alone. Now we’ll start writing the test cases.

Before we do that, we need to import the go-mock package:

`1`	`go get github.com/xhd2015/go-mock`

Then add the test code bar_test.go:

package bar

import (
  "context"
  "fmt"
  "testing"
  
  // Importing generated helper packages
  mock_foo "hello/test/mock_gen/foo"
)

func TestHi(t *testing.T) {
  ctx := context.Background()
  // Setting the Mock Functions
  ctx = mock_foo.Setup(ctx, func(m *mock_foo.M) {
    // Each function has a corresponding variable
    m.Hi = func(ctx context.Context, a, b int) int {
      return a - b
    }
  })
  // At this point, calling Hi() will call the function you just mocked.
  fmt.Println(Hi(ctx, 1, 2))
}

You can’t run go test directly, you have to use go-mock instead.

`1`	`go-mock test hello/bar`

You will see that the output is 1 - 2 + 1 = 0 instead of the original 1 + 2 + 1 = 4. Mock succeeded.

Now let’s summarise the advantages of go-mock.

First of all, it is a pure Go language implementation, almost no black technology, no platform requirements, and no dependency on Go ABI, so it has the best compatibility, supports all platforms, and supports almost all Go language versions. Because of the source code level dependency, the subsequent new version of Go language will not have any impact on go-mock. This is the best thing about go-mock!

Second, it has no concurrency security issues, because all Mock state is passed through a ctx, which is the officially recommended standard practice.

Based on these two points, I judge that go-mock is a potential stock with a promising future. And I hope to be involved in this project in the future.

But is go-mock perfect? Of course not! And the drawbacks are just as obvious. These drawbacks are just what I found through my limited testing and research, or maybe I don’t understand them, or they can be solved later. Here is a list for your reference.

Firstly, it can only Mock code written by itself. Functions like time.Now() can’t be mocked, and as far as I know, it doesn’t seem to be able to mock functions in third-party packages.

Second, the function to be mocked must be passed a ctx argument. This isn’t a big problem, since most functions in the Go ecosystem now support passing ctx. However, there are still a lot of scenarios where you don’t have to pass ctx. It’s a bit of a pain in the neck to have to pass ctx to be able to Mock. It’s also because of this dependency on ctx that it doesn’t currently support mocking functions that return ctx. This is a small price to pay for concurrency safety.

Third, scanning and generating code may slow down the test speed.

Fourth, Mock generic functions are not supported. Again, this is not a big problem.

Fifth, the current implementation does not support Mock functions in the current package, which creates circular dependencies. But it should be solved.

In general go-mock has a very good future. Because business projects most likely will not directly call standard functions or third-party packages, but will provide a layer of encapsulation. In this case, the shortcomings of go-mock can be bypassed. Code at different levels of the project can be Mocked at will, which already solves most of the problems. Using the 2-8 principle, go-mock is indeed a very good tool.

My personal suggestion is to try both options: Monkey Patch has some risk, but is more flexible; go-mock has no risk at all, but there are some limitations in the mock scenarios. For the time being, they can’t be replaced by each other, so it’s up to you to choose according to your own business scenarios and team realities.

At this point in the writing the main point of this article should be over. However, the beginning of the article mentioned “early programme” actually did not have a suitable place to place. For the sake of completeness, I’ll put it at the end, as a supplement to the whole article for your reference.

Early programme

As I said earlier, because in 2018 I lacked sufficient knowledge of the Go language itself, let alone a hacky solution like Monkey Patch. And I couldn’t afford to implement goroutine isolation at that time, so I couldn’t turn on parallel testing if I used it.

Under the conditions at that time, we chose an earthy and solid solution - functional simulation.

The so-called functional simulation is to make up for what is missing. If the code needs a database, we run a MySQL instance; if the code needs caching, we run a Redis instance; if the code needs to call external APIs, we use jarcoal/httpmock to Mock, and in order to keep the different pipelines from interfering with each other, we use Docker mode to run the pipelines. To keep the unit test environment consistent with the code, we automatically import table structures from the development database before running the use cases. To generate test data, we fill the database with data via a special seed.sql.

This system achieves only one goal, which is that the unit tests of different Merge Requests can run in parallel without interfering with each other. But Pipelines in the same branch could only run serially because they shared a common data persistence layer. Over the course of five years this led to our test cases getting slower and slower. Another obvious problem was that most of the test cases were in the DAO layer, which is the closest to the data and the easiest to write. The higher the layer, the more code dependencies there are, the harder the use cases are to write, and naturally there are fewer and fewer of them.

After the team split up, when the other groups went on to make the new service, everyone actually used the same solution. I wouldn’t have known about it if it weren’t for the problem of not being able to export table structures in the middle of an IT environment adjustment. The Mock solution I shared with you earlier was almost ignored, and the newer solutions in the industry are even more unappreciated. When the old solution is slow, the first thing that comes to mind is to improve on the old solution. One improvement is to mark different table structure dependencies for different packages, and then create different databases for different packages before startup, so as to achieve data isolation, and then open concurrency testing.

Ref: https://taoshu.in/go/mock.html

Table of Contents

Code Dependencies

Dependency Injection

Monkey Patch

Code rewrite

Early programme