使用pgxmock进行Golang PostgreSQL单元测试

pgxmock是一个用于pgx数据库驱动的模拟库，它允许你在不连接真实数据库的情况下进行单元测试。下面我将详细介绍如何使用pgxmock。

安装pgxmock

首先安装pgxmock库：

go get github.com/pashagolub/pgxmock

基本用法

1. 创建模拟连接

import (
	"testing"
	"github.com/pashagolub/pgxmock"
)

func TestSomething(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatalf("an error '%s' was not expected when opening a stub database connection", err)
	}
	defer mock.Close(context.Background())
	
	// 在这里设置期望并测试你的代码
}

2. 设置查询期望

// 设置期望查询并返回模拟行
mock.ExpectQuery("SELECT name FROM users WHERE id = \\$1").
	WithArgs(1).
	WillReturnRows(pgxmock.NewRows([]string{"name"}).AddRow("John"))

// 执行测试代码
var name string
err := mock.QueryRow(context.Background(), "SELECT name FROM users WHERE id = $1", 1).Scan(&name)
if err != nil {
	t.Errorf("error was not expected: %s", err)
}

if name != "John" {
	t.Errorf("expected name to be 'John', got '%s'", name)
}

// 确保所有期望都满足
if err := mock.ExpectationsWereMet(); err != nil {
	t.Errorf("there were unfulfilled expectations: %s", err)
}

3. 设置执行期望

// 设置期望执行
mock.ExpectExec("UPDATE users SET name = \\$1 WHERE id = \\$2").
	WithArgs("Jane", 1).
	WillReturnResult(pgxmock.NewResult("UPDATE", 1))

// 执行测试代码
result, err := mock.Exec(context.Background(), "UPDATE users SET name = $1 WHERE id = $2", "Jane", 1)
if err != nil {
	t.Errorf("error was not expected: %s", err)
}

affected := result.RowsAffected()
if affected != 1 {
	t.Errorf("expected 1 affected row, got %d", affected)
}

完整示例

下面是一个完整的存储库层单元测试示例：

package repository

import (
	"context"
	"testing"
	
	"github.com/jackc/pgx/v4"
	"github.com/pashagolub/pgxmock"
	"github.com/stretchr/testify/assert"
)

type UserRepository struct {
	conn *pgx.Conn
}

func NewUserRepository(conn *pgx.Conn) *UserRepository {
	return &UserRepository{conn: conn}
}

func (r *UserRepository) GetUserName(id int) (string, error) {
	var name string
	err := r.conn.QueryRow(context.Background(), "SELECT name FROM users WHERE id = $1", id).Scan(&name)
	return name, err
}

func TestUserRepository_GetUserName(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatalf("failed to create mock: %v", err)
	}
	defer mock.Close(context.Background())

	// 设置期望
	mock.ExpectQuery("SELECT name FROM users WHERE id = \\$1").
		WithArgs(1).
		WillReturnRows(pgxmock.NewRows([]string{"name"}).AddRow("John Doe"))

	repo := NewUserRepository(mock)
	name, err := repo.GetUserName(1)

	assert.NoError(t, err)
	assert.Equal(t, "John Doe", name)

	// 确保所有期望都满足
	assert.NoError(t, mock.ExpectationsWereMet())
}

高级功能

1. 事务测试

func TestTransaction(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatalf("failed to create mock: %v", err)
	}
	defer mock.Close(context.Background())

	// 期望开始事务
	mock.ExpectBegin()
	
	// 期望在事务中的查询
	mock.ExpectQuery("SELECT name FROM users WHERE id = \\$1").
		WithArgs(1).
		WillReturnRows(pgxmock.NewRows([]string{"name"}).AddRow("John"))
	
	// 期望提交事务
	mock.ExpectCommit()

	// 执行测试代码
	tx, err := mock.Begin(context.Background())
	if err != nil {
		t.Fatalf("failed to begin transaction: %v", err)
	}

	var name string
	err = tx.QueryRow(context.Background(), "SELECT name FROM users WHERE id = $1", 1).Scan(&name)
	if err != nil {
		tx.Rollback(context.Background())
		t.Fatalf("failed to query: %v", err)
	}

	err = tx.Commit(context.Background())
	if err != nil {
		t.Fatalf("failed to commit: %v", err)
	}

	assert.NoError(t, mock.ExpectationsWereMet())
}

2. 错误模拟

func TestErrorCase(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatalf("failed to create mock: %v", err)
	}
	defer mock.Close(context.Background())

	// 模拟查询错误
	mock.ExpectQuery("SELECT name FROM users WHERE id = \\$1").
		WithArgs(1).
		WillReturnError(fmt.Errorf("some error"))

	// 执行测试代码
	var name string
	err = mock.QueryRow(context.Background(), "SELECT name FROM users WHERE id = $1", 1).Scan(&name)
	
	assert.Error(t, err)
	assert.Equal(t, "some error", err.Error())
	assert.NoError(t, mock.ExpectationsWereMet())
}

最佳实践

1. 每个测试用例创建新的mock：避免测试间的相互影响
2. 总是检查ExpectationsWereMet：确保所有预期的查询/执行都被调用
3. 使用正则表达式匹配查询：使用\\$而不是$来匹配参数占位符
4. 清理资源：记得调用Close()或使用defer
5. 结合测试框架：如testify/assert可以简化断言代码

pgxmock是测试pgx数据库交互的强大工具，它可以帮助你编写快速、隔离的单元测试，而无需依赖真实的数据库连接。