如何测试gRPC服务器流式响应？Go+Cobra客户端测试求助

我正在用Go开发一个基于gRPC的客户端，客户端用Cobra CLI构建，用来替换项目旧CLI接口，同时作为未来的API客户端。目前已经完成2个基础命令的编写和测试，这类命令只需发送请求、接收响应后关闭连接。现在我在开发一个从服务器获取流式响应的命令，客户端功能已实现，但只能对接真实服务器，希望编写测试用例方便后续扩展功能。

服务器Protobuf定义 #

message StreamEventsRequest {
    repeated string policies = 1;
    google.protobuf.FieldMask mask = 2;
}

message StreamEventsResponse {
    Event event = 1;
}

service TraceeService {
    rpc StreamEvents(StreamEventsRequest) returns (stream StreamEventsResponse);
}

streamEvents命令代码 #

package cmd

import (
    "github.com/ShohamBit/TraceeClient/client"
    pb "github.com/aquasecurity/tracee/api/v1beta1"
    "github.com/spf13/cobra"
)

var (
    streamEventsCmd = &cobra.Command{
        Use:   "streamEvents",
        Short: "Stream events from tracee",
        Long:  "Stream the events that tracee trace to the client",
        Run: func(cmd *cobra.Command, args []string) {
            streamEvents(cmd, args)
        },
    }
)

func streamEvents(cmd *cobra.Command, args []string) {
    // create service client
    client, err := client.NewServiceClient(serverInfo)
    if err != nil {
        cmd.PrintErrln("Error creating client: ", err)
    }
    defer client.CloseConnection()
    // stream events
    req := &pb.StreamEventsRequest{Policies: args}
    stream, err := client.StreamEvents(cmd.Context(), req)
    if err != nil {
        cmd.PrintErrln("Error calling StreamEvents: ", err)
    }
    // Receive and process streamed responses
    for {
        res, err := stream.Recv()
        if err != nil {
            cmd.PrintErrln("Error receiving streamed event: ", err)
            break
        }
        cmd.Println(res.Event)
    }

}

我尝试编写测试用例和Mock服务器，但没达到预期效果，希望获得编写gRPC Mock流的方法。

解决方案：用testify/mock和gRPC Mock生成工具实现流式Mock #

步骤1：安装依赖

先安装gRPC Mock生成工具和testify测试库：

go install github.com/golang/mock/mockgen@latest
go get github.com/stretchr/testify/mock

步骤2：抽象客户端接口并生成Mock代码

首先在client包中抽象出客户端接口（如果还没有），方便后续Mock：

// client/client.go
import (
    "context"
    pb "github.com/aquasecurity/tracee/api/v1beta1"
    "google.golang.org/grpc"
)

type TraceeServiceClient interface {
    StreamEvents(ctx context.Context, req *pb.StreamEventsRequest, opts ...grpc.CallOption) (pb.TraceeService_StreamEventsClient, error)
    CloseConnection() error
}

// 确保原NewServiceClient返回这个接口类型
func NewServiceClient(serverInfo ServerInfo) (TraceeServiceClient, error) {
    // 原实现代码...
}

然后运行mockgen生成Mock客户端：

mockgen -source=client/client.go -destination=client/mock_client.go -package=client

步骤3：实现流式响应Mock

编写自定义的流式客户端Mock，模拟Recv()方法返回预设的事件流：

package cmd

import (
    "context"
    "io"
    "testing"

    "github.com/ShohamBit/TraceeClient/client"
    pb "github.com/aquasecurity/tracee/api/v1beta1"
    "github.com/golang/mock/gomock"
    "github.com/spf13/cobra"
    "github.com/stretchr/testify/assert"
)

// 自定义流式响应Mock，实现pb.TraceeService_StreamEventsClient接口
type mockStreamClient struct {
    events []*pb.StreamEventsResponse
    index  int
}

func (m *mockStreamClient) Recv() (*pb.StreamEventsResponse, error) {
    if m.index >= len(m.events) {
        // 返回EOF表示流结束
        return nil, io.EOF
    }
    res := m.events[m.index]
    m.index++
    return res, nil
}

// 接口要求的其他空实现
func (m *mockStreamClient) SendMsg(msg interface{}) error { return nil }
func (m *mockStreamClient) RecvMsg(msg interface{}) error { return nil }

步骤4：编写测试用例

替换原代码的客户端创建逻辑，注入Mock客户端并验证流程：

// 先修改cmd包的代码，支持依赖注入
// cmd/stream_events.go
var newClientFunc = client.NewServiceClient // 新增全局变量，用于测试替换

func streamEvents(cmd *cobra.Command, args []string) {
    client, err := newClientFunc(serverInfo) // 改用注入的函数创建客户端
    if err != nil {
        cmd.PrintErrln("Error creating client: ", err)
        return // 原代码缺少错误返回，建议补充
    }
    defer client.CloseConnection()

    req := &pb.StreamEventsRequest{Policies: args}
    stream, err := client.StreamEvents(cmd.Context(), req)
    if err != nil {
        cmd.PrintErrln("Error calling StreamEvents: ", err)
        return
    }

    for {
        res, err := stream.Recv()
        if err != nil {
            cmd.PrintErrln("Error receiving streamed event: ", err)
            break
        }
        cmd.Println(res.Event)
    }
}

// 测试用例
func TestStreamEvents(t *testing.T) {
    ctrl := gomock.NewController(t)
    defer ctrl.Finish()

    // 创建Mock客户端
    mockClient := client.NewMockTraceeServiceClient(ctrl)
    // 构造测试用的流式响应
    testEvents := &mockStreamClient{
        events: []*pb.StreamEventsResponse{
            {Event: &pb.Event{Id: "1", Name: "test-event-1"}},
            {Event: &pb.Event{Id: "2", Name: "test-event-2"}},
        },
    }

    // 设置Mock预期：调用StreamEvents时返回预设的流和无错误
    mockClient.EXPECT().StreamEvents(
        gomock.Any(),
        &pb.StreamEventsRequest{Policies: []string{"test-policy"}},
    ).Return(testEvents, nil)
    mockClient.EXPECT().CloseConnection().Return(nil)

    // 替换客户端创建函数为Mock实现
    originalNewClient := newClientFunc
    newClientFunc = func(_ client.ServerInfo) (client.TraceeServiceClient, error) {
        return mockClient, nil
    }
    defer func() { newClientFunc = originalNewClient }()

    // 准备测试命令
    cmd := streamEventsCmd
    cmd.SetArgs([]string{"test-policy"})
    cmd.SetContext(context.Background())

    // 执行命令并验证无错误
    err := cmd.Execute()
    assert.NoError(t, err)
}

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