Go中NATS订阅测试：解决*nats.Subscription无法Mock的问题

核心思路是依赖抽象而非具体实现，通过定义抽象接口隔离NATS的具体类型，这样就能在测试中轻松Mock这些接口。

步骤1：定义抽象接口 #

把业务代码中用到的NATS核心能力抽成接口，包括订阅上下文、订阅实例、消息处理三个层面：

// 抽象NATS连接上下文，包含订阅方法
type NatsContext interface {
    SubscribeSync(subject string, opts ...nats.SubOpt) (NatsSubscription, error)
}

// 抽象订阅实例，包含消息获取和取消订阅方法
type NatsSubscription interface {
    NextMsgWithContext(ctx context.Context) (NatsMsg, error)
    Unsubscribe() error
}

// 抽象NATS消息，包含业务代码中用到的操作方法
type NatsMsg interface {
    Data() []byte
    InProgress()
    Ack()
    Nak()
    Term()
}

步骤2：修改业务代码依赖抽象 #

调整JetStreamConnection结构体，将原来的具体*nats.Conn类型替换为NatsContext接口：

type JetStreamConnection struct {
    context      NatsContext // 替换原*nats.Conn类型
    unmarshaller func(data []byte, v interface{}) error
    // 其他原有字段...
}

同时修改SubscribeMultiple中涉及消息的代码，将*nats.Msg替换为NatsMsg接口：

func (conn *JetStreamConnection) SubscribeMultiple(ctx context.Context, subject string,
    subscribers ...*SubscriptionCallback) error {

    callbacks := make(map[string]func(*pnats.NatsMessage) (func(context.Context), error))
    for _, subscriber := range subscribers {
        callbacks[subscriber.Category] = subscriber.Callback
    }

    fullSubject := fmt.Sprintf("%s.*", subject)
    sub, err := conn.context.SubscribeSync(fullSubject, nats.Context(ctx))
    if err != nil {
        return err
    }

loop:
    for {
        select {
        case <-ctx.Done():
            break loop
        default:
        }

        msg, err := sub.NextMsgWithContext(ctx)
        if err != nil {
            return err
        }

        msg.InProgress()

        var message pnats.NatsMessage
        if err := conn.unmarshaller(msg.Data(), &message); err != nil {
            msg.Term()
            return err
        }

        subscriber, ok := callbacks[message.Context.Category]
        if !ok {
            msg.Nak()
            continue
        }

        callback, err := subscriber(&message)
        if err == nil {
            msg.Ack()
        } else {
            msg.Nak()
            return err
        }

        callback(ctx)
    }

    if err := sub.Unsubscribe(); err != nil {
        return err
    }

    return nil
}

步骤3：实现Mock对象用于测试 #

手动编写Mock实现（或用testify/mock简化），控制返回值并记录方法调用：

// MockNatsMsg 模拟NATS消息，记录关键方法调用
type MockNatsMsg struct {
    DataBytes        []byte
    InProgressCalled bool
    AckCalled        bool
    NakCalled        bool
    TermCalled       bool
}

func (m *MockNatsMsg) Data() []byte                { return m.DataBytes }
func (m *MockNatsMsg) InProgress()                 { m.InProgressCalled = true }
func (m *MockNatsMsg) Ack()                        { m.AckCalled = true }
func (m *MockNatsMsg) Nak()                        { m.NakCalled = true }
func (m *MockNatsMsg) Term()                       { m.TermCalled = true }

// MockNatsSubscription 模拟订阅实例
type MockNatsSubscription struct {
    NextMsgFunc     func(ctx context.Context) (NatsMsg, error)
    UnsubscribeFunc func() error
}

func (m *MockNatsSubscription) NextMsgWithContext(ctx context.Context) (NatsMsg, error) {
    if m.NextMsgFunc != nil {
        return m.NextMsgFunc(ctx)
    }
    return nil, nil
}

func (m *MockNatsSubscription) Unsubscribe() error {
    if m.UnsubscribeFunc != nil {
        return m.UnsubscribeFunc()
    }
    return nil
}

// MockNatsContext 模拟NATS连接上下文
type MockNatsContext struct {
    SubscribeSyncFunc func(subject string, opts ...nats.SubOpt) (NatsSubscription, error)
}

func (m *MockNatsContext) SubscribeSync(subject string, opts ...nats.SubOpt) (NatsSubscription, error) {
    if m.SubscribeSyncFunc != nil {
        return m.SubscribeSyncFunc(subject, opts...)
    }
    return nil, nil
}

步骤4：编写测试用例 #

注入Mock对象，验证业务逻辑的正确性：

func TestSubscribeMultiple_RouteSuccess(t *testing.T) {
    // 准备测试消息数据
    testMsgData := []byte(`{"Context":{"Category":"test"},"Payload":"demo"}`)
    mockMsg := &MockNatsMsg{DataBytes: testMsgData}

    // 模拟订阅：返回测试消息，超时后触发ctx.Done退出循环
    mockSub := &MockNatsSubscription{
        NextMsgFunc: func(ctx context.Context) (NatsMsg, error) {
            select {
            case <-ctx.Done():
                return nil, ctx.Err()
            default:
                return mockMsg, nil
            }
        },
        UnsubscribeFunc: func() error { return nil },
    }

    // 模拟NATS上下文：验证订阅主题是否正确
    mockCtx := &MockNatsContext{
        SubscribeSyncFunc: func(subject string, opts ...nats.SubOpt) (NatsSubscription, error) {
            if subject != "test.*" {
                return nil, fmt.Errorf("invalid subject: %s", subject)
            }
            return mockSub, nil
        },
    }

    // 准备测试回调：标记是否被执行
    callbackExecuted := false
    testSubscriber := &SubscriptionCallback{
        Category: "test",
        Callback: func(msg *pnats.NatsMessage) (func(context.Context), error) {
            callbackExecuted = true
            return func(ctx context.Context) {}, nil
        },
    }

    // 创建测试用的JetStreamConnection实例
    conn := &JetStreamConnection{
        context:      mockCtx,
        unmarshaller: json.Unmarshal,
    }

    // 带超时的上下文，确保测试能自动结束
    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()

    // 执行测试方法
    err := conn.SubscribeMultiple(ctx, "test", testSubscriber)

    // 验证结果
    if err != nil {
        t.Fatalf("SubscribeMultiple failed: %v", err)
    }
    if !callbackExecuted {
        t.Error("expected subscriber callback to be executed")
    }
    if !mockMsg.InProgressCalled {
        t.Error("expected msg.InProgress() to be called")
    }
    if !mockMsg.AckCalled {
        t.Error("expected msg.Ack() to be called on success")
    }
}

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