如何为含内部函数调用的Golang函数编写单元测试及Mock MongoDB客户端

一、修复mtest创建Mock Client返回nil的问题 #

你直接在mtest.New()后获取mt.Client会得到nil，因为Mock Client需要在mt.Run()的测试用例上下文里才会被初始化。正确的使用方式是把测试逻辑放在mt.Run()的闭包中：

func TestSaveProjectInfo(t *testing.T) {
    // 初始化mtest实例
    mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
    defer mt.Close()

    // 在Run方法的回调中执行测试逻辑
    mt.Run("successful save", func(mt *mtest.T) {
        // 此时mt.Client已被正确初始化，不会是nil
        client := mt.Client

        // 后续构造测试数据、调用函数等逻辑
    })
}

二、Mock内部调用的UpsertProjectInfo函数 #

要MockUpsertProjectInfo，需要先解耦代码，让依赖可替换。这里提供两种方案：

方案1：通过接口依赖注入（推荐） #

1. 定义一个包含UpsertProjectInfo方法的接口：

type ProjectUpserter interface {
    UpsertProjectInfo(col *mongo.Collection, info processor.ProjectInfo) error
}

2. 为原UpsertProjectInfo函数实现这个接口：

type DefaultUpserter struct{}

func (u *DefaultUpserter) UpsertProjectInfo(col *mongo.Collection, info processor.ProjectInfo) error {
    // 这里放原UpsertProjectInfo的实现代码
}

3. 修改SaveProjectInfo函数，增加接口参数：

func SaveProjectInfo(projectInfo processor.ProjectInfo, configMap map[string]string, client *mongo.Client, upserter ProjectUpserter) error {
    collection := client.Database(configMap[constants.Dbname_str_key]).Collection(configMap[constants.Db_collection_str])
    return upserter.UpsertProjectInfo(collection, projectInfo)
}

4. 使用testify mock创建Mock实现并编写测试：

import "github.com/stretchr/testify/mock"

// 定义Mock结构体
type MockUpserter struct {
    mock.Mock
}

// 实现ProjectUpserter接口
func (m *MockUpserter) UpsertProjectInfo(col *mongo.Collection, info processor.ProjectInfo) error {
    args := m.Called(col, info)
    return args.Error(0)
}

// 测试用例
func TestSaveProjectInfo(t *testing.T) {
    mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
    defer mt.Close()

    mt.Run("success case", func(mt *mtest.T) {
        // 构造测试数据
        testInfo := processor.ProjectInfo{/* 填充测试字段 */}
        testConfig := map[string]string{
            constants.Dbname_str_key: "test_db",
            constants.Db_collection_str: "test_col",
        }

        // 创建MockUpserter并设置预期
        mockUpserter := new(MockUpserter)
        // 预期调用UpsertProjectInfo，传入任意Collection和testInfo，返回nil错误
        mockUpserter.On("UpsertProjectInfo", mock.Anything, testInfo).Return(nil)

        // 调用SaveProjectInfo
        err := SaveProjectInfo(testInfo, testConfig, mt.Client, mockUpserter)

        // 断言结果
        assert.NoError(t, err)
        // 验证Mock的预期是否被触发
        mockUpserter.AssertExpectations(t)
    })
}

方案2：替换包级变量（快速修改） #

如果不想修改函数参数，可以把UpsertProjectInfo定义为包级变量，测试时替换为Mock函数：

1. 修改原代码：

// 定义包级变量，指向原UpsertProjectInfo函数
var upsertProjectInfo = UpsertProjectInfo

func SaveProjectInfo(projectInfo processor.ProjectInfo, configMap map[string]string, client *mongo.Client) error {
    collection := client.Database(configMap[constants.Dbname_str_key]).Collection(configMap[constants.Db_collection_str])
    // 调用包级变量
    return upsertProjectInfo(collection, projectInfo)
}

2. 编写测试：

func TestSaveProjectInfo(t *testing.T) {
    mt := mtest.New(t, mtest.NewOptions().ClientType(mtest.Mock))
    defer mt.Close()

    mt.Run("success case", func(mt *mtest.T) {
        // 保存原函数，测试后恢复
        originalUpsert := upsertProjectInfo
        defer func() { upsertProjectInfo = originalUpsert }()

        testInfo := processor.ProjectInfo{/* 填充测试字段 */}
        testConfig := map[string]string{
            constants.Dbname_str_key: "test_db",
            constants.Db_collection_str: "test_col",
        }

        // 标记Mock是否被调用
        mockCalled := false
        // 替换为Mock函数
        upsertProjectInfo = func(col *mongo.Collection, info processor.ProjectInfo) error {
            mockCalled = true
            // 断言参数是否符合预期
            assert.Equal(t, "test_db", col.Database().Name())
            assert.Equal(t, "test_col", col.Name())
            assert.Equal(t, testInfo, info)
            return nil
        }

        // 调用函数
        err := SaveProjectInfo(testInfo, testConfig, mt.Client)

        // 断言结果
        assert.NoError(t, err)
        assert.True(t, mockCalled)
    })
}

内容的提问来源于stack exchange，提问作者Pramit Pakhira