Flexible mocking for testing in Go

SafetyCulture’s codebase is made up of many micro-services, many of which are written in Go. These micro-services talk to other micro-services and may connect to one or more datastores.

Source: https://github.com/marcusolsson/gophers/blob/master/gopherdata-gopher.png

Mocking is an essential tool to help unit test our code paths without using a real [integration] client.

In Go (and other languages) it’s a very understood pattern to use interfaces to achieve mocking.

TL;DR

Packages should return concrete (usually pointer or struct) types: that way, new methods can be added to implementations without requiring extensive refactoring.

We still take advantage of Go’s implicit interfaces to aid in testing/mocking, and it can all be done without a testing/mocking framework.

Key takeaways:

• avoid large interfaces
• only use the interfaces/methods that you need
• use nil and empty structs to mock your happy path
• customize your mocks with function properties

The Anti-Pattern

Throughout our codebase it’s very common to see code like this:

type Store interface {
    ListUsers() ([]*User, error)
    ListTeamsByUser(userID int) ([]*Team, error)
    ListAddressesByUser(userID int) (*[]Address, error)
    // ...etc
    GetUser(id int) (*User, error)
    GetTeam(id int) (*Team, error)
    // ...etc
    CreateUser(user *User) (int, error)
    CreateTeam(userID int, team *Team) (int, error)
    // ...etc
    SetUserDefaultAddress(userID int, addressID int) error
    // ...etc
}
type store struct {
    db sql.DB
}
func NewStore(conn string) Store {
    db := initDB(conn)
    return &store{db:db}
}

What we observe above is an interface with a large number of methods and a constructor that returns the Store interface.

But we know that this is an anti-pattern (we’ll be addressing this later):

Do not define interfaces on the implementor side of an API “for mocking”; instead, design the API so that it can be tested using the public API of the real implementation.

The Good

Ignoring that it’s an anti-pattern for now; at first glance, the above pattern appears to make sense and we can easily use our Store interface wherever we need to by passing the same interface:

SomeHandler(ctx context.Context, store db.Store, userID int) error {
    user, err := store.GetUser(userID)
    if err != nil {
        return err
    }
    // ...
    team, err := store.GetTeam(user.DefaultTeam)
    if err != nil {
        return err
    }
    // ...
}
SomeOtherHandler(ctx context.Context, store db.Store, firstname, lastname string) error {
    user := User{Firstname: firstname, Lastname: lastname}
    if err := store.CreateUser(&user); err != nil {
        return err
    }
    // ...
}

The Bad

This “convenience” comes at a cost; the hidden problem is that in order to test the two functions above we need to mock the entire Store interface, even when we are only using a sub set of the functions (some method bodies ignored for brevity):

type storeMock struct {
    getUser *User
    getUserErr error
    getTeam *Team
    getTeamErr error
    // ...
}
func (m *storeMock) ListUsers() ([]*User, error) {}
func (m *storeMock) ListTeamsByUser(userID int) ([]*Team, error) {}
func (m *storeMock) ListAddressesByUser(userID int) (*[]Address, error) {}
func (m *storeMock) GetUser(id int) (*User, error) {
    return m.getUser, m.getUserErr
}
func (m *storeMock) GetTeam(id int) (*Team, error) {
    return m.getTeam, m.getTeamErr
}
func (m *storeMock) CreateUser(user *User) (int, error) {}
func (m *storeMock) CreateTeam(userID int, team *Team) (int, error) {}
func (m *storeMock) SetUserDefaultAddress(userID int, addressID int) error {}

Creating a mock for each method we want to test would be very painful, and hard to maintain, so we tend to use a single mock struct for all our tests.

Every time we add a new interface method we would need to update our mocks and maybe even have to update our tests.

The Ugly

Many of our teams turned to auto-generation of these mocks that work in tandem with testing frameworks.

This solves the problem of maintaining our own test mocks, but can still lead to “hard to maintain” tests.

Lets step through an example to discover some of the problems that we may encounter:

// handler.go
SomeHandler(ctx context.Context, store db.Store, userID int) error {
    user, err := store.GetUser(userID)
    if err != nil {
        return err
    }
    // ...
    team, err := store.GetTeam(user.DefaultTeam)
    if err != nil {
        return err
    }
    // ...
}
// handler_test.go
func TestSomeHandler(t *testing.T) {
    tests := [...]struct {
        name string
        store func() db.Store
        shouldErr bool
    }{
        {
            "HappyPath",
            func() db.Store {
                mock := &dbmock.Store{}
                mock.On("GetUser", mock.Anything).Return(&db.User{}, nil)
                mock.On("GetTeam", mock.Anything).Return(&db.Team{}, nil)
                return mock
            },
            false,
        },
 }
 for _, tt := range tests {
     tt := tt
     t.Run(tt.name, func(t *testing.T){
         err := SomeHandler(context.Background(), tt.store(), 0)
         if tt.shouldErr && err == nil {
              t.Error("expected error but got <nil>")
         }
      })
 }

The above test is a typical example (simplified) of a test that uses mocking. In this example we are only testing the “Happy Path” that skips over the error paths, but we can add more test, and using a similar pattern, test all code paths by mocking the input/output of the mocks.

Problem One

We add a method to our interface:

type Store interface {
    ListUsers() ([]*User, error)
    ListTeamsByUser(userID int) ([]*Team, error)
    
    // ...etc
    AddUserToTeam(teamID, userID int) error
}

Running our tests will fail with an error similar to:

FAIL: cannot use *dbmock.Store as db.Store; missing method AddUserToTeam(teamID, userID int) error

If we have auto-generated our mocks this is a simple fix, we just need to re-generate, but if we created the mock manually, we would need to add this function.

Problem Two

We now need to add this method to the body of our SomeHandler :

SomeHandler(ctx context.Context, store db.Store, userID int) error {
    user, err := store.GetUser(userID)
    // ...
    err := store.AddUserToTeam(ateamID, userID)
    if err != nil {
        return err
    }
    team, err := store.GetTeam(user.DefaultTeam)
    // ...
}

Running our tests again we get something like this error:

mock: I don't know what to return because the method call was unexpected.
Either do Mock.On("AddUserToTeam").Return(...) first, or remove the AddUserToTeam() call.
This method was unexpected

We can’t remove the call, so I guess we’ll need to mock the method in our test. But here is the kicker: we’ll have to update ALL our tests where the code path reaches this function, and if you have lots of tests (which you should) this gets repetitive:

{
    "HappyPath",
    func() db.Store {
        mock := &dbmock.Store{}
        mock.On("GetUser", mock.Anything).Return(&db.User{}, nil)
        
        mock.On("AddUserToTeam", mock.Anything, mock.Anything).Return(nil)
        
        mock.On("GetTeam", mock.Anything).Return(&db.Team{}, nil)
        return mock
    },
    false,
},

Problem Three

What happens if we refactor our Store interface? Thanks to static typing and the compiler, refactoring in Go is normally simple (very contrived example, but you should get the point):

type Store interface {
    // ...etc
    --- GetTeam(id int) (*Team, error)
    +++ GetSquad(id int) (*Team, error)
   // ...etc
}

This results in a similar error as before:

mock: I don't know what to return because the method call was unexpected.
Either do Mock.On("GetSquad").Return(...) first, or remove the GetSquad() call.
This method was unexpected

Because this mocking framework is using strings we can’t capitalise on Go refactoring tools, but have to find/replace "GetTeam" with the new method signature.

Real Implementation

Lets start by removing our large Store interface and returning the real implementation of our store so that we can start “accepting interfaces, return structs” and move away from our anti-pattern:

type Store struct {
    db sql.DB
}
func NewStore(conn string) *Store {
    db := initDB(conn)
    return &store{db:db}
}

We still need to use interfaces to aid in mocking; interfaces should describe the behavior of your functions or methods required, thus reducing the coupling between the behavior and the implementation of that behavior.

Loose coupling between your packages ultimately results in more maintainable and readable code.

So rather than creating a large interface we create small interfaces (1 or 2 methods) for only the functions we need and compose interfaces if we need more than 1.

type UserGetter interface {
    GetUser(id int) (*db.User, error)
}
type TeamGetter interface {
    GetTeam(id int) (*db.Team, error)
}
type UserTeamGetter interface {
    UserGetter
    TeamGetter
}
type UserCreator interface {
    CreateUser(user *db.User) (int, error)
}
SomeHandler(ctx context.Context, store UserTeamGetter, userID int) error {
    user, err := store.GetUser(userID)
    if err != nil {
        return err
    }
    // ...
    team, err := store.GetTeam(user.DefaultTeam)
    if err != nil {
        return err
    }
    // ...
}
SomeOtherHandler(ctx context.Context, store UserCreator, firstname, lastname string) error {
    user := User{Firstname: firstname, Lastname: lastname}
    if err := store.CreateUser(&user); err != nil {
        return err
    }
    // ...
}

With the functions only using the interfaces they need make it easier to test and we don’t need to change the caller code that uses the real store implementation.

store := db.NewStore(conn)
SomeHandler(ctx, store, userID)

Happy Path

In many test cases we need to have a general happy path through our mocked functions, so to reduce repetitive code we can make the nil struct (or empty struct ie. &myMock{}) as the mock for the happy path:

type userTeamMock struct {}
func (*userTeamMock) GetUser(id int) (*db.User, error) {
    return &db.User{Firstname: "foo", Lastname: "bar"}, nil
}
func (*userTeamMock) GetTeam(id int) (*db.Team, error) {
    return &db.Team{Name: "foobar"}, nil
}
func TestSomeHandler(t *testing.T) {
    tests := [...]struct {
        name string
        store *userTeamMock
        shouldErr bool
    }{
        {
            "HappyPath",
            nil,
            false,
        },
 }
 for _, tt := range tests {
     tt := tt
     t.Run(tt.name, func(t *testing.T){
         err := SomeHandler(context.Background(), tt.store, 0)
         if tt.shouldErr && err == nil {
              t.Error("expected error but got <nil>")
         }
      })
 }

Mock Functions

The reason for creating a mock is that we can test more code branches than just the happy path. We can customize our mock to return anything we want to mock using function properties.

type userTeamMock struct {
    getUserFn func(id int) (*db.User, error)
    getTeamFn func(id int) (*db.Team, error)
}
func (m *userTeamMock) GetUser(id int) (*db.User, error) {
    if m != nil && m.getUserFn != nil {
        return m.getUserFn(id)
    }
    return &db.User{Firstname: "foo", Lastname: "bar"}, nil
}
func (*userTeamMock) GetTeam(id int) (*db.Team, error) {
    if m != nil && m.getTeamFn != nil {
        return m.getTeamFn(id)
    }
    return &db.Team{Name: "foobar"}, nil
}

Using my mock now becomes much simpler; my happy path will always work if the userTeamMock is nil or empty and I can override that whenever we need:

{
    "GetUserError",
    &userTeamMock{
        getUserFn: func GetUser(id int) (*db.User, error) {
           return nil, errors.New("user error")
        }
    },
    true,
},
{
    "GetTeamError",
    &userTeamMock{
        getTeamFn: func GetTeam(id int) (*db.Team, error) {
           return nil, errors.New("team error")
        }
    },
    true,
},
{
    "HappyPath",
    nil,
    false,
},

With our new mock, there is no need to auto-generate our mocks or use complex testing frameworks, the Go programing language makes testing and mocking easy straight out of the box.

Problems Be Gone

Now that we are returning the concrete Store pointer type, any new methods that are added do not require large refactoring of our test as they are only using the interfaces that they require. Problem one solved.

If we add a method that we need in a function we’ll have to add it to our interface:

type UserTeamGetAdder interface {
    UserTeamGetter
    AddUserToTeam(teamID, userID int) error
}
SomeHandler(ctx context.Context, store UserTeamGetAdder, userID int) error {
    // ...
}

This will case our tests to fail, but the Go compiler will lead you the way:

cannot use &(store literal) (value of type *userTeamMock) as UserTeamGetAdder value in argument to SomeHandler: missing method AddUserToTeam

Rather than refactoring all of our test cases we can just add the method to our mock:

func (m *userTeamMock) AddUserToTeam(teamID, userID int) (*db.User, error) {
    if m != nil && m.getUserFn != nil {
        return m.addUserToTeamFn(teamID, userID)
    }
    return &db.User{Firstname: "foo", Lastname: "bar"}, nil
}

All our test will now pass without changing any of our test code, we’re only left with adding more test cases. Problem 2 solved.

We’ve also avoided any magic strings which makes any refactoring required a lot easier by utilizing compiler errors (and other Go tools like gopls LSP) rather than runtime errors. Problem 3 solved.

Final Thoughts

Some of the reasons we enjoy the Go Programming Language is that it’s easy to learn, concise, maintainable and very readable.

But sometimes we still make it hard for ourselves as software engineers with many design patterns we’ve inherited along the way.

Go’s simplicity challenges us to think differently; and questions whether we really need framework X or framework Y; maybe the standard library that Go provides will be enough?

“Simplicity is complicated”
- Rob Pike