如何使用stretchr/testify mock AWS SDK的Client类方法？

问题根源 #

1. 方法签名不匹配：你编写的Mock方法PostToConnection()缺少真实方法的上下文、参数和可选函数参数，Go会将其判定为不同方法，导致调用时仍走真实Client的实现。
2. 依赖具体类型而非抽象接口：业务代码直接依赖*apigatewaymanagementapi.Client具体类型，而非SDK提供的抽象接口，无法替换为Mock实现。

1. 修改业务代码依赖抽象接口 #

AWS SDK v2为每个客户端提供了对应的API接口（如apigatewaymanagementapi.ClientAPI），该接口定义了客户端的所有方法。将业务函数的参数改为依赖此接口：

import "github.com/aws/aws-sdk-go-v2/service/apigatewaymanagementapi"

func publishData(ctx context.Context, api apigatewaymanagementapi.ClientAPI) {
    input := &apigatewaymanagementapi.PostToConnectionInput{
        ConnectionId: ..., // 业务所需参数
        Data:          ...,
    }
    _, err := api.PostToConnection(ctx, input)
    // 原有业务逻辑保持不变
}

2. 创建匹配签名的Mock结构体 #

基于testify/mock实现ClientAPI接口的PostToConnection方法，确保签名与真实接口完全一致：

import (
    "context"
    "github.com/aws/aws-sdk-go-v2/service/apigatewaymanagementapi"
    "github.com/stretchr/testify/mock"
)

type MockAPIGatewayManagementAPI struct {
    mock.Mock
}

// 严格匹配ClientAPI接口的PostToConnection方法签名
func (m *MockAPIGatewayManagementAPI) PostToConnection(
    ctx context.Context,
    params *apigatewaymanagementapi.PostToConnectionInput,
    optFns ...func(*apigatewaymanagementapi.Options),
) (*apigatewaymanagementapi.PostToConnectionOutput, error) {
    args := m.Called(ctx, params, optFns)
    return args.Get(0).(*apigatewaymanagementapi.PostToConnectionOutput), args.Error(1)
}

// 若业务代码用到ClientAPI的其他方法，需按同样方式实现对应的Mock方法

3. 编写测试用例 #

初始化Mock、设置调用预期并验证：

import (
    "context"
    "testing"
    "github.com/aws/aws-sdk-go-v2/service/apigatewaymanagementapi"
    "github.com/stretchr/testify/assert"
    "github.com/stretchr/testify/mock"
)

func TestPublishData(t *testing.T) {
    // 初始化Mock实例
    mockAPI := new(MockAPIGatewayManagementAPI)

    // 构造预期的输入参数
    expectedInput := &apigatewaymanagementapi.PostToConnectionInput{
        ConnectionId: aws.String("test-conn-id"),
        Data:          []byte("test-data"),
    }

    // 设置方法调用预期：匹配任意上下文、指定输入参数、任意可选函数，返回预设结果
    mockAPI.On("PostToConnection", mock.Anything, expectedInput, mock.Anything).
        Return(&apigatewaymanagementapi.PostToConnectionOutput{}, nil)

    // 调用业务函数
    publishData(context.Background(), mockAPI)

    // 验证Mock的方法是否按预期被调用
    mockAPI.AssertExpectations(t)
}

关键注意事项 #

• 必须保证Mock方法的签名与ClientAPI接口完全一致，包括参数类型、数量和可变参数。
• 业务代码依赖抽象接口而非具体实现，是实现Mock的核心，符合依赖倒置原则。
• 若业务仅用到ClientAPI的部分方法，Mock只需实现被用到的方法即可，未实现的方法被调用会触发panic。

内容的提问来源于stack exchange，提问作者Joey Yi Zhao