使用pgxmock进行Golang PostgreSQL单元测试

pgxmock是一个用于测试PostgreSQL数据库交互的模拟库，特别适合与pgx库一起使用。它允许你在不连接真实数据库的情况下测试数据库相关代码。

安装pgxmock

首先安装pgxmock库：

go get github.com/pashagolub/pgxmock

基本用法

1. 创建模拟连接

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

func TestDatabaseOperation(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. 设置查询期望

func TestGetUser(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close(context.Background())

	// 设置期望查询和返回的行
	rows := mock.NewRows([]string{"id", "name", "email"}).
		AddRow(1, "John Doe", "john@example.com").
		AddRow(2, "Jane Doe", "jane@example.com")

	mock.ExpectQuery("SELECT id, name, email FROM users WHERE id = \\$1").
		WithArgs(1).
		WillReturnRows(rows)

	// 调用实际函数
	user, err := GetUser(mock, 1)
	if err != nil {
		t.Errorf("error was not expected while getting user: %s", err)
	}

	// 验证期望
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expectations: %s", err)
	}
	
	// 验证返回的数据
	if user.Name != "John Doe" {
		t.Errorf("expected name to be John Doe, got %s", user.Name)
	}
}

3. 测试事务

func TestTransaction(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close(context.Background())

	// 期望事务开始
	mock.ExpectBegin()
	
	// 期望查询
	mock.ExpectQuery("SELECT balance FROM accounts WHERE id = \\$1").
		WithArgs(1).
		WillReturnRows(mock.NewRows([]string{"balance"}).AddRow(100))
	
	// 期望更新
	mock.ExpectExec("UPDATE accounts SET balance = balance - \\$1 WHERE id = \\$2").
		WithArgs(50, 1).
		WillReturnResult(pgxmock.NewResult("UPDATE", 1))
	
	// 期望事务提交
	mock.ExpectCommit()

	// 调用实际函数
	err = TransferMoney(mock, 1, 2, 50)
	if err != nil {
		t.Errorf("error was not expected while transferring money: %s", err)
	}

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

4. 测试错误场景

func TestQueryError(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close(context.Background())

	// 设置期望查询返回错误
	expectedErr := errors.New("connection failed")
	mock.ExpectQuery("SELECT \\* FROM users").
		WillReturnError(expectedErr)

	// 调用实际函数
	_, err = GetAllUsers(mock)
	if err == nil {
		t.Error("expected error but got none")
	}
	
	if err != expectedErr {
		t.Errorf("expected error '%v' but got '%v'", expectedErr, err)
	}

	// 验证期望
	if err := mock.ExpectationsWereMet(); err != nil {
		t.Errorf("there were unfulfilled expectations: %s", err)
	}
}

高级用法

1. 使用匹配器

func TestWithMatcher(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close(context.Background())

	// 使用AnyArg匹配任意参数
	mock.ExpectExec("UPDATE users SET name = \\$1 WHERE id = \\$2").
		WithArgs(pgxmock.AnyArg(), 1).
		WillReturnResult(pgxmock.NewResult("UPDATE", 1))

	// 调用实际函数
	err = UpdateUserName(mock, 1, "New Name")
	if err != nil {
		t.Errorf("error was not expected while updating user: %s", err)
	}
}

2. 测试预处理语句

func TestPreparedStatement(t *testing.T) {
	mock, err := pgxmock.NewConn()
	if err != nil {
		t.Fatal(err)
	}
	defer mock.Close(context.Background())

	// 期望预处理语句
	mock.ExpectPrepare("get_user", "SELECT name FROM users WHERE id = \\$1")
	
	// 期望执行预处理语句
	rows := mock.NewRows([]string{"name"}).AddRow("John Doe")
	mock.ExpectQuery("get_user").
		WithArgs(1).
		WillReturnRows(rows)

	// 调用实际函数
	name, err := GetUserName(mock, 1)
	if err != nil {
		t.Errorf("error was not expected while getting user name: %s", err)
	}
	
	if name != "John Doe" {
		t.Errorf("expected name to be John Doe, got %s", name)
	}
}

最佳实践

1. 每个测试用例创建新的模拟连接：避免测试之间的状态污染
2. 总是检查ExpectationsWereMet：确保所有期望的查询都被执行
3. 使用defer关闭连接：防止资源泄漏
4. 测试错误场景：不仅要测试成功路径，还要测试错误处理
5. 保持测试独立：每个测试应该只关注一个功能点

pgxmock是一个强大的工具，可以让你在不依赖真实数据库的情况下彻底测试数据库交互代码。通过合理设置期望和验证结果，你可以确保你的数据库层代码在各种情况下都能正确工作。