如何Mock未通过接口声明的第三方包函数 以confluent-kafka-go为例

核心逻辑是通过自行定义抽象接口做依赖倒置，解除业务代码和第三方包具体实现的耦合，不需要修改第三方包源码即可实现Mock。

实现步骤 #

1. 自定义消费者抽象接口 #

仅声明业务中实际用到的方法即可：

import "github.com/confluentinc/confluent-kafka-go/kafka"

// KafkaConsumer 自定义消费者接口
type KafkaConsumer interface {
    SubscribeTopics(topics []string, rebalanceCb kafka.RebalanceCb) error
    Poll(timeoutMs int) kafka.Event
    Close() error
}

2. 实现原生消费者的适配层 #

封装官方的*kafka.Consumer结构体，实现上述自定义接口：

// RealKafkaConsumer 真实Kafka消费者适配实现
type RealKafkaConsumer struct {
    inner *kafka.Consumer
}

// NewRealKafkaConsumer 真实消费者构造函数
func NewRealKafkaConsumer(conf *kafka.ConfigMap) (KafkaConsumer, error) {
    c, err := kafka.NewConsumer(conf)
    if err != nil {
        return nil, err
    }
    return &RealKafkaConsumer{inner: c}, nil
}

func (r *RealKafkaConsumer) SubscribeTopics(topics []string, rebalanceCb kafka.RebalanceCb) error {
    return r.inner.SubscribeTopics(topics, rebalanceCb)
}

func (r *RealKafkaConsumer) Poll(timeoutMs int) kafka.Event {
    return r.inner.Poll(timeoutMs)
}

func (r *RealKafkaConsumer) Close() error {
    return r.inner.Close()
}

3. 业务代码依赖抽象接口 #

将业务逻辑中所有直接依赖*kafka.Consumer的地方替换为依赖自定义的KafkaConsumer接口：

// 示例业务服务
type MessageProcessService struct {
    consumer KafkaConsumer
    topic    string
}

func NewMessageProcessService(consumer KafkaConsumer, topic string) *MessageProcessService {
    return &MessageProcessService{consumer: consumer, topic: topic}
}

// 业务方法示例
func (s *MessageProcessService) Run() error {
    if err := s.consumer.SubscribeTopics([]string{s.topic}, nil); err != nil {
        return err
    }
    for {
        ev := s.consumer.Poll(100)
        // 业务处理逻辑省略
        _ = ev
    }
}

4. 编写Mock实现并执行单元测试 #

可以手动写Mock，也可以用gomock等工具自动生成Mock代码，以下是手动Mock的测试示例：

// MockKafkaConsumer 测试用Mock实现
type MockKafkaConsumer struct {
    MockSubscribeTopics func(topics []string, rebalanceCb kafka.RebalanceCb) error
    MockPoll            func(timeoutMs int) kafka.Event
    MockClose           func() error
}

func (m *MockKafkaConsumer) SubscribeTopics(topics []string, rebalanceCb kafka.RebalanceCb) error {
    return m.MockSubscribeTopics(topics, rebalanceCb)
}

func (m *MockKafkaConsumer) Poll(timeoutMs int) kafka.Event {
    return m.MockPoll(timeoutMs)
}

func (m *MockKafkaConsumer) Close() error {
    return m.MockClose()
}

// 单元测试示例
func TestMessageProcessService_Run(t *testing.T) {
    // 预设返回的测试消息
    expectMsg := &kafka.Message{
        Value: []byte("test payload"),
        TopicPartition: kafka.TopicPartition{
            Topic:  &[]string{"test-topic"}[0],
            Offset: 10,
        },
    }

    // 初始化Mock实例，预设方法返回值
    mockCsm := &MockKafkaConsumer{
        MockSubscribeTopics: func(topics []string, _ kafka.RebalanceCb) error {
            if len(topics) != 1 || topics[0] != "test-topic" {
                t.Errorf("unexpected subscribe topics: %v", topics)
            }
            return nil
        },
        MockPoll: func(_ int) kafka.Event {
            return expectMsg
        },
        MockClose: func() error {
            return nil
        },
    }

    // 传入Mock实例初始化业务服务
    svc := NewMessageProcessService(mockCsm, "test-topic")
    // 执行业务方法、断言结果，省略后续逻辑
    _ = svc
}

注意事项 #

• 自定义接口不需要覆盖原生消费者的全部方法，仅声明业务用到的方法即可，减少冗余代码
• 如果使用gomock工具，可直接在接口上方加//go:generate mockgen -source=consumer.go -destination=mock_consumer.go -package=kafka注释，执行go generate即可自动生成Mock代码，不需要手动编写Mock结构体
• 仅初始化消费者的位置需要替换为自定义的NewRealKafkaConsumer构造函数，其余业务逻辑无需修改，侵入性极低

内容的提问来源于stack exchange，提问作者Krishna patil