在这篇博客文章中，我们将探讨如何在 Go 语言中使用 RedisMock 库和 SQLMock 库进行单元测试。我们将通过一个简单的示例项目展示如何为数据库操作编写单元测试，并确保我们的代码按预期工作。
项目结构
我们的项目包含以下文件：
-
sql.go：包含数据库操作的实现。 -
sql_test.go：包含数据库操作的单元测试。 -
redis.go：包含 Redis 操作的实现。 -
redis_test.go：包含 Redis 操作的单元测试。
数据库操作实现 sql.go
首先，我们定义了一个 User 结构体和一个 Data 结构体，用于封装数据库操作。
package main
import (
"context"
"time"
"gorm.io/gorm"
)
type User struct {
ID int `gorm:"primaryKey column:id"`
Birthday time.Time `gorm:"column:birthday"`
Username string `gorm:"column:username"`
}
func (u User) TableName() string {
return "users"
}
type Data struct {
DB *gorm.DB
}
func (d *Data) Create(ctx context.Context, user *User) error {
return d.DB.Create(user).Error
}
func (d *Data) Get(ctx context.Context, id int) (*User, error) {
var user User
err := d.DB.Where("`id` = ?", id).First(&user).Error
return &user, err
}
func (d *Data) Update(ctx context.Context, user *User) error {
return d.DB.Save(user).Error
}
func (d *Data) Delete(ctx context.Context, id int) error {
return d.DB.Where(`id = ?`, id).Delete(User{}).Error
}
数据库操作单元测试 sql_test.go
接下来，我们使用 SQLMock 库为数据库操作编写单元测试。
package main
import (
"context"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/stretchr/testify/assert"
"gorm.io/driver/mysql"
"gorm.io/gorm"
)
func initTestDB() (*gorm.DB, sqlmock.Sqlmock, error) {
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(sqlmock.QueryMatcherEqual))
if err != nil {
return nil, nil, err
}
mock.ExpectQuery("SELECT VERSION()").WillReturnRows(sqlmock.NewRows([]string{"version"}).AddRow("8.0.23"))
gormDB, err := gorm.Open(mysql.New(mysql.Config{
Conn: db,
}))
if err != nil {
return nil, nil, err
}
return gormDB, mock, nil
}
func TestData_Create(t *testing.T) {
gormDB, mock, err := initTestDB()
assert.NoError(t, err)
birthday, err := time.Parse("2006-01-02", "2021-01-01")
assert.NoError(t, err)
user := &User{
ID: 1,
Birthday: birthday,
Username: "test",
}
mock.ExpectBegin()
mock.ExpectExec("INSERT INTO `users` (`birthday`,`username`,`id`) VALUES (?,?,?)").
WithArgs(user.Birthday, user.Username, user.ID).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectCommit()
data := &Data{DB: gormDB}
err = data.Create(context.Background(), user)
assert.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestData_Get(t *testing.T) {
gormDB, mock, err := initTestDB()
assert.NoError(t, err)
userID := 1
birthday, err := time.Parse("2006-01-02", "2021-01-01")
assert.NoError(t, err)
rows := sqlmock.NewRows([]string{"id", "birthday", "username"}).
AddRow(userID, birthday, "test")
mock.ExpectQuery("SELECT * FROM `users` WHERE `id` = ? ORDER BY `users`.`id` LIMIT ?").
WithArgs(userID, 1).
WillReturnRows(rows)
data := &Data{DB: gormDB}
user, err := data.Get(context.Background(), userID)
assert.NoError(t, err)
assert.Equal(t, userID, user.ID)
assert.Equal(t, birthday, user.Birthday)
assert.Equal(t, "test", user.Username)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestData_Update(t *testing.T) {
gormDB, mock, err := initTestDB()
assert.NoError(t, err)
birthday, err := time.Parse("2006-01-02", "2021-01-01")
assert.NoError(t, err)
user := &User{
ID: 1,
Birthday: birthday,
Username: "test",
}
mock.ExpectBegin()
mock.ExpectExec("UPDATE `users` SET `birthday`=?,`username`=? WHERE `id` = ?").
WithArgs(user.Birthday, user.Username, user.ID).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectCommit()
data := &Data{DB: gormDB}
err = data.Update(context.Background(), user)
assert.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestData_Delete(t *testing.T) {
gormDB, mock, err := initTestDB()
assert.NoError(t, err)
userID := 1
mock.ExpectBegin()
mock.ExpectExec("DELETE FROM `users` WHERE id = ?").
WithArgs(userID).
WillReturnResult(sqlmock.NewResult(1, 1))
mock.ExpectCommit()
data := &Data{DB: gormDB}
err = data.Delete(context.Background(), userID)
assert.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
Redis 操作实现 redis.go
我们还定义了一个 SessionStore 结构体，用于封装 Redis 操作。
package main
import (
"context"
"time"
"github.com/redis/go-redis/v9"
)
type SessionStore struct {
redisClient redis.UniversalClient
}
func NewSessionStore(redisClient redis.UniversalClient) *SessionStore {
return &SessionStore{redisClient: redisClient}
}
func (s *SessionStore) GetSession(ctx context.Context, key string) (string, error) {
return s.redisClient.Get(ctx, key).Result()
}
func (s *SessionStore) SetSession(ctx context.Context, key, value string, expiration time.Duration) error {
return s.redisClient.Set(ctx, key, value, expiration).Err()
}
Redis 操作单元测试 redis_test.go
最后，我们使用 Redismock 库为 Redis 操作编写单元测试。
import (
"context"
"testing"
"time"
"github.com/go-redis/redismock/v9"
"github.com/stretchr/testify/assert"
)
func TestSessionStore_GetSession(t *testing.T) {
db, mock := redismock.NewClientMock()
store := NewSessionStore(db)
key := "test-key"
value := "test-value"
mock.ExpectGet(key).SetVal(value)
result, err := store.GetSession(context.Background(), key)
assert.Nil(t, err)
assert.Equal(t, value, result)
assert.NoError(t, mock.ExpectationsWereMet())
}
func TestSessionStore_SetSession(t *testing.T) {
db, mock := redismock.NewClientMock()
store := NewSessionStore(db)
key := "test-key"
value := "test-value"
expiration := 10 * time.Minute
mock.ExpectSet(key, value, expiration).SetVal("OK")
err := store.SetSession(context.Background(), key, value, expiration)
assert.NoError(t, err)
assert.NoError(t, mock.ExpectationsWereMet())
}
执行单元测试
通过go test -v ./...对当前项目进行单元测试。结果如下：
image.png
总结
通过本文，我们展示了如何使用 SQLMock 和 Redismock 库为 Go 语言中的数据库和 Redis 操作编写单元测试。希望这篇文章能帮助你更好地理解和应用这些工具，提高代码的可靠性和可维护性。
