Golang 单元测试 - 数据层

前言

今天我们先来看看有关数据层(repo)的单元测试应该如何实践。

数据层，就是我们常常说的 repo/dao，其功能就是和数据库、缓存或者其他数据源打交道。它需要从数据源中获取数据，并返回给上一层。在这一层通常没有复杂业务的逻辑，所以最重要的就是测试各个数据字段的编写是否正确，以及 SQL 等查询条件是否正常能被筛选。

当然，数据层也基本上是最底层了，通常这一层的单元测试更加的重要，因为如果一个字段名称和数据库不一致上层所有依赖这个方法的地方全部都会报错。

由于数据层和数据源打交道，那么测试的麻烦点就在于，通常我们不能要求外接一定能提供一个数据源供我们测试：一方面是由于我们不可能随时都能连上测试服务器的数据库，另一方面我们也不能要求单元测试运行的时候只有你一个人在使用这个数据库，而且数据库数据干净。退一步讲，我们也没办法 mock，如果 mock 了 sql，那么测试的意义就不大了。

下面我们就以我们常见的 mysql 数据库为例，看看在 golang 中如何进行单元测试的编写。

准备工作的说明

数据源

首先，我们需要一个干净的数据源，由于我们没有办法依赖于外部服务器的数据库，那么我们就利用最常使用的 docker 来帮助我们构建一个所需要使用的数据源。

我们这里使用 github.com/ory/dockertest 来帮助我们构建测试的环境，它能帮助我们启动一个所需要的环境，当然你也可以选择手动使用 docker 或者 docker-compose 来创建。

初始数据

有了数据库之后，我们还需要表结构和初始数据，这部分也有两种方案：

使用 orm 提供的 sync/migration 类似的功能，将结构体直接映射生成表字段，通过 go 代码创建初始数据
直接使用 sql 语句，通过执行 sql 语句来创建对应的表结构和字段数据
本案例使用第一种方式进行，第二种也类似

基本 case 代码

我们首先来快速搞定一下默认的 case 代码，也就是我们常常搬砖的 CRUD。(这里仅给出最基本的实现，重点主要关注在单元测试上)

package repo

import (
	"context"

	"go-demo/m/unit-test/entity"
	"xorm.io/xorm"
)

type UserRepo interface {
	AddUser(ctx context.Context, user *entity.User) (err error)
	DelUser(ctx context.Context, userID int) (err error)
	GetUser(ctx context.Context, userID int) (user *entity.User, exist bool, err error)
}

type userRepo struct {
	db *xorm.Engine
}

func NewUserRepo(db *xorm.Engine) UserRepo {
	return &userRepo{db: db}
}

func (ur userRepo) AddUser(ctx context.Context, user *entity.User) error {
	_, err := ur.db.Insert(user)
	return err
}

func (ur userRepo) DelUser(ctx context.Context, userID int) error {
	_, err := ur.db.Delete(&entity.User{ID: userID})
	return err
}

func (ur userRepo) GetUser(ctx context.Context, userID int) (user *entity.User, exist bool, err error) {
	user = &entity.User{ID: userID}
	exist, err = ur.db.Get(user)
	return user, exist, err
}

初始化测试环境

首先创建 repo_main_test.go 文件

package repo

import (
	"database/sql"
	"fmt"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/ory/dockertest/v3"
	"github.com/ory/dockertest/v3/docker"
	"go-demo/m/unit-test/entity"
	"xorm.io/xorm"
	"xorm.io/xorm/schemas"
)

type TestDBSetting struct {
	Driver       string
	ImageName    string
	ImageVersion string
	ENV          []string
	PortID       string
	Connection   string
}

var (
	mysqlDBSetting = TestDBSetting{
		Driver:       string(schemas.MYSQL),
		ImageName:    "mariadb",
		ImageVersion: "10.4.7",
		ENV:          []string{"MYSQL_ROOT_PASSWORD=root", "MYSQL_DATABASE=linkinstar", "MYSQL_ROOT_HOST=%"},
		PortID:       "3306/tcp",
		Connection:   "root:root@(localhost:%s)/linkinstar?parseTime=true",
	}
	tearDown       func()
	testDataSource *xorm.Engine
)

func TestMain(t *testing.M) {
	defer func() {
		if tearDown != nil {
			tearDown()
		}
	}()
	if err := initTestDataSource(mysqlDBSetting); err != nil {
		panic(err)
	}
	if ret := t.Run(); ret != 0 {
		panic(ret)
	}
}

func initTestDataSource(dbSetting TestDBSetting) (err error) {
	connection, imageCleanUp, err := initDatabaseImage(dbSetting)
	if err != nil {
		return err
	}
	dbSetting.Connection = connection

	testDataSource, err = initDatabase(dbSetting)
	if err != nil {
		return err
	}

	tearDown = func() {
		testDataSource.Close()
		imageCleanUp()
	}
	return nil
}

func initDatabaseImage(dbSetting TestDBSetting) (connection string, cleanup func(), err error) {
	pool, err := dockertest.NewPool("")
	pool.MaxWait = time.Minute * 5
	if err != nil {
		return "", nil, fmt.Errorf("could not connect to docker: %s", err)
	}

	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
		Repository: dbSetting.ImageName,
		Tag:        dbSetting.ImageVersion,
		Env:        dbSetting.ENV,
	}, func(config *docker.HostConfig) {
		config.AutoRemove = true
		config.RestartPolicy = docker.RestartPolicy{Name: "no"}
	})
	if err != nil {
		return "", nil, fmt.Errorf("could not pull resource: %s", err)
	}

	connection = fmt.Sprintf(dbSetting.Connection, resource.GetPort(dbSetting.PortID))
	if err := pool.Retry(func() error {
		db, err := sql.Open(dbSetting.Driver, connection)
		if err != nil {
			fmt.Println(err)
			return err
		}
		return db.Ping()
	}); err != nil {
		return "", nil, fmt.Errorf("could not connect to database: %s", err)
	}
	return connection, func() { _ = pool.Purge(resource) }, nil
}

func initDatabase(dbSetting TestDBSetting) (dbEngine *xorm.Engine, err error) {
	dbEngine, err = xorm.NewEngine(dbSetting.Driver, dbSetting.Connection)
	if err != nil {
		return nil, err
	}
	err = initDatabaseData(dbEngine)
	if err != nil {
		return nil, fmt.Errorf("init database data failed: %s", err)
	}
	return dbEngine, nil
}

func initDatabaseData(dbEngine *xorm.Engine) error {
	return dbEngine.Sync(new(entity.User))
}

下面说明其中的方法和要点

TestMain：这个方法是 golang test 的一个特性，它会在所有单元测试之前自动执行，特别适合用于初始化数据和清理测试遗留环境。这个方法中 tearDown 是为了清理连接和镜像用的
initDatabaseImage：方法主要就是利用 github.com/ory/dockertest 提供功能拉取一个对应的 docker 镜像并启动
initDatabaseData：方法主要利用了 xorm 的 Sync 方法去初始化了数据库，当然这里也可以构建你所需要的初始化数据，比如你需要初始化一个超级管理员等等
testDataSource：我们将最终初始化的数据源放在了这里，由于后续单元测试的时候使用，这里由于只有一个数据库，就没有封装

编写单元测试

有了前面的准备工作，单元测试就变得简单了

package repo

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"go-demo/m/unit-test/entity"
)

func Test_userRepo_AddUser(t *testing.T) {
	ur := NewUserRepo(testDataSource)
	user := &entity.User{
		Username: "LinkinStar",
	}
	err := ur.AddUser(context.TODO(), user)
	assert.NoError(t, err)

	dbUser, exist, err := ur.GetUser(context.TODO(), user.ID)
	assert.NoError(t, err)
	assert.True(t, exist)
	assert.Equal(t, user.Username, dbUser.Username)

	err = ur.DelUser(context.TODO(), user.ID)
	assert.NoError(t, err)
}

可以看到我们只需要像平常写代码一样直接调用对应的方法就可以进行单元测试了。

其中 https://github.com/stretchr/testify 是一个非常好用的断言工具，能帮助我们快速实现单元测试中的断言，以便我们快速确定单元测试是否正确。
单元测试需要注意的是，我们这里测试的是添加用户，也就是插入数据，为保证单元测试的独立性，测试完当前方法后数据应该保持一致，故需要进行数据删除，以保证不会干扰到其他的单元测试。

注意事项

本地需要有 docker 环境
第一次启动由于需要拉取镜像，根据网络情况不同，拉取时间不同
正常情况下，我们设定了 AutoRemove 为 true 并且不再重启，测试完成之后会将测试使用的 mysql 镜像关闭并删除，但是如果测试意外中断，或者强制中断时，会导致镜像被遗留下来。故，本地测试之后可以使用 docker ps 命令查看是否有遗留
当然根据所需要的数据源不同，你可以使用其他不同的镜像进行操作，效果是一样的