一文搞懂Go gRPC服務Handler單元測試
在云原生時代和微服務架構背景下，HTTP和RPC協議成為服務間通信和與客戶端交互的兩種主要方式。對于Go語言而言，標準庫提供了net/http/httptest包，為開發人員提供了便捷的方式來構建服務端HTTP Handler單元測試的測試腳手架代碼，而無需真正建立HTTP服務器，讓開發人員可以聚焦于對Handler業務邏輯的測試。比如下面這個示例：
// grpc-test-examples/httptest/http_handler_test.go
func myHandler(w http.ResponseWriter, r *http.Request) {
// 設置響應頭
w.Header().Set("Content-Type", "text/plain")
// 根據請求方法進行不同的處理
switch r.Method {
case http.MethodGet:
// 處理GET請求
fmt.Fprint(w, "Hello, World!")
... ...
}
}
func TestMyHandler(t *testing.T) {
// 創建一個ResponseRecorder來記錄Handler的響應
rr := httptest.NewRecorder()
// 創建一個模擬的HTTP請求，可以指定請求的方法、路徑、正文等
req, err := http.NewRequest("GET", "/path", nil)
if err != nil {
t.Fatal(err)
}
// 調用被測試的Handler函數，傳入ResponseRecorder和Request對象
// 這里假設被測試的Handler函數為myHandler
myHandler(rr, req)
// 檢查響應狀態碼和內容
if rr.Code != http.StatusOK {
t.Errorf("Expected status 200; got %d", rr.Code)
}
expected := "Hello, World!"
if rr.Body.String() != expected {
t.Errorf("Expected body to be %q; got %q", expected, rr.Body.String())
}
}注：對http client端的單元測試，也可以利用httptest的NewServer來構建一個fake的http server[1]。
然而，對于使用主流的gRPC等RPC協議的服務端Handler[2]來說，是否存在類似httptest的測試腳手架生成工具包呢？對gRPC的服務端Handler有哪些單元測試的方法呢？在這篇文章中，我們就一起來探究一下。
1. 建立被測的gRPC服務端Handler
我們首先來建立一個涵蓋多種gRPC通信模式的服務端Handler集合。
gRPC支持四種通信模式，它們分別為：
- 簡單RPC(Simple RPC，也稱為Unary RPC)
這是最簡單的，也是最常用的gRPC通信模式，簡單來說就是一請求一應答。
- 服務端流RPC(Server-streaming RPC)
客戶端發來一個請求，服務端通過流返回多個應答。
- 客戶端流RPC(Client-streaming RPC)
客戶端通過流發來多個請求，服務端以一個應答回復。
- 雙向流RPC(Bidirectional-Streaming RPC)
客戶端通過流發起多個請求，服務端也通過流對應返回多個應答。
注：關于gRPC四種通信方式的詳情，可以參考我之前寫的《gRPC客戶端的那些事兒[3]》一文。
我們這個SUT(被測目標)是包含以上四種通信模式的gRPC服務，它的Protocol Buffers[4]文件如下：
// grpc-test-examples/grpctest/IDL/proto/mygrpc.proto
syntax = "proto3";
package mygrpc;
service MyService {
// Unary RPC
rpc UnaryRPC(RequestMessage) returns (ResponseMessage) {}
// Server-Streaming RPC
rpc ServerStreamingRPC(RequestMessage) returns (stream ResponseMessage) {}
// Client-Streaming RPC
rpc ClientStreamingRPC(stream RequestMessage) returns (ResponseMessage) {}
// Bidirectional-Streaming RPC
rpc BidirectionalStreamingRPC(stream RequestMessage) returns (stream ResponseMessage) {}
}
message RequestMessage {
string message = 1;
}
message ResponseMessage {
string message = 1;
}通過protoc，我們可基于上述proto文件生成MyService樁(Stub)代碼，生成的代碼放在了mygrpc目錄下面：
// grpc-test-examples/grpctest/Makefile
all: gen
gen:
protoc -I ./IDL/proto mygrpc.proto --gofast_out=plugins=grpc:./mygrpc注：你的環境下需要安裝protoc[5]和protoc-gen-go[6]才能正確執行上面生成命令，具體的安裝方法可參考protoc安裝文檔[7]。
注：除了使用經典的protoc[8]基于proto文件生成Go源碼外，也可以基于Go開發的buf cli[9]進行代碼生成和API管理。buf cLi是現代、快速、高效的Protobuf API管理的終極工具，為基于Protobuf的開發和維護提供了全面的解決方案。等有機會的時候，我在以后的文章中詳細說說buf。
有了生成的樁代碼后，我們便可以建立一個gRPC服務器：
// grpc-test-examples/grpctest/main.go
package main
import (
pb "demo/mygrpc"
"log"
"net"
"google.golang.org/grpc"
)
func main() {
// 創建 gRPC 服務器
lis, err := net.Listen("tcp", ":50051")
if err != nil {
log.Fatalf("failed to listen: %v", err)
}
s := grpc.NewServer()
// 注冊 MyService 服務
pb.RegisterMyServiceServer(s, &server{})
// 啟動 gRPC 服務器
log.Println("Starting gRPC server...")
if err := s.Serve(lis); err != nil {
log.Fatalf("failed to serve: %v", err)
}
}我們看到：在main函數中，我們創建了一個TCP監聽器，并使用grpc.NewServer()創建了一個gRPC服務器。然后，我們通過調用pb.RegisterMyServiceServer()將server類型的實例注冊到gRPC服務器上，以處理來自客戶端的請求。最后，我們啟動gRPC服務器并監聽指定的端口。
上面代碼中注冊到服務器中的server類型就是實現了MyService服務接口的具體類型，它實現了MyService定義的所有方法：
// grpc-test-examples/grpctest/server.go
package main
import (
"context"
"fmt"
"strconv"
pb "demo/mygrpc"
)
type server struct{}
func (s *server) UnaryRPC(ctx context.Context, req *pb.RequestMessage) (*pb.ResponseMessage, error) {
message := "Unary RPC received: " + req.Message
fmt.Println(message)
return &pb.ResponseMessage{
Message: "Unary RPC response",
}, nil
}
func (s *server) ServerStreamingRPC(req *pb.RequestMessage, stream pb.MyService_ServerStreamingRPCServer) error {
message := "Server Streaming RPC received: " + req.Message
fmt.Println(message)
for i := 0; i < 5; i++ {
response := &pb.ResponseMessage{
Message: "Server Streaming RPC response " + strconv.Itoa(i+1),
}
if err := stream.Send(response); err != nil {
return err
}
}
return nil
}
func (s *server) ClientStreamingRPC(stream pb.MyService_ClientStreamingRPCServer) error {
var messages []string
for {
req, err := stream.Recv()
if err != nil {
return err
}
messages = append(messages, req.Message)
if req.Message == "end" {
break
}
}
message := "Client Streaming RPC received: " + fmt.Sprintf("%v", messages)
fmt.Println(message)
return stream.SendAndClose(&pb.ResponseMessage{
Message: "Client Streaming RPC response",
})
}
func (s *server) BidirectionalStreamingRPC(stream pb.MyService_BidirectionalStreamingRPCServer) error {
for {
req, err := stream.Recv()
if err != nil {
return err
}
message := "Bidirectional Streaming RPC received: " + req.Message
fmt.Println(message)
response := &pb.ResponseMessage{
Message: "Bidirectional Streaming RPC response",
}
if err := stream.Send(response); err != nil {
return err
}
}
}在上面代碼中，我們創建了一個server結構體類型，并實現了MyService的所有RPC方法。每個方法都接收相應的請求消息，并返回對應的響應消息。我們的目標僅是演示如何對上述gRPC Handler進行單元測試，所以這里的實現邏輯非常簡單。
接下來，我們就來逐一對這些gRPC的Handler方法進行單測，我們先從簡單的UnaryRPC方法開始。
2. Unary RPC Handler的單元測試
Unary RPC是最簡單，也是最容易理解的RPC通信模式，即客戶端與服務端采用一請求一應答的模式。server類型的UnaryRPC Handler方法的原型如下：
// grpc-test-examples/grpctest/server.go
func (s *server) UnaryRPC(ctx context.Context, req *pb.RequestMessage) (*pb.ResponseMessage, error)就像文章開頭做的那個httpserver的handler單測一樣，我們肯定不想真實啟動一個gRPC server，也不想測試gRPC服務器本身。我們只想測試服務端handler方法的邏輯是否正確。
觀察一下這個方法原型，我們發現它僅依賴兩個消息結構：RequestMessage和ResponseMessage，這兩個消息結構是上面基于proto文件自動生成的，這樣我們就可以不借助任何工具包實現對UnaryRPC handler方法的單測，也無需啟動真實的gRPC Server：
// grpc-test-examples/grpctest/server_test.go
type server struct{}
func TestServerUnaryRPC(t *testing.T) {
s := &server{}
req := &pb.RequestMessage{
Message: "Test message",
}
resp, err := s.UnaryRPC(context.Background(), req)
if err != nil {
t.Fatalf("UnaryRPC failed: %v", err)
}
expectedResp := &pb.ResponseMessage{
Message: "Unary RPC response",
}
if resp.Message != expectedResp.Message {
t.Errorf("Unexpected response. Got: %s, Want: %s", resp.Message, expectedResp.Message)
}
}將其改造為基于subtest[10]和表驅動的測試也非常easy：
// grpc-test-examples/grpctest/server_test.go
func TestServerUnaryRPCs(t *testing.T) {
tests := []struct {
name string
requestMessage *pb.RequestMessage
expectedResp *pb.ResponseMessage
}{
{
name: "Test Case 1",
requestMessage: &pb.RequestMessage{
Message: "Test message",
},
expectedResp: &pb.ResponseMessage{
Message: "Unary RPC response",
},
},
// Add more test cases as needed
}
s := &server{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resp, err := s.UnaryRPC(context.Background(), tt.requestMessage)
if err != nil {
t.Fatalf("UnaryRPC failed: %v", err)
}
if resp.Message != tt.expectedResp.Message {
t.Errorf("Unexpected response. Got: %s, Want: %s", resp.Message, tt.expectedResp.Message)
}
})
}
}如果gRPC handler測試都像UnaryRPC這樣簡單那就好了，但實際上...，好吧，我們繼續向下看就好了。
3. 針對Streaming通信模式的單元測試
3.1 ServerStreamingRPC的測試
前面說過，gRPC支持三種Streaming通信模式：Server-Streaming RPC、Client-Streaming RPC和Bidirectional-Streaming RPC。
我們先來看看Server-Streaming RPC的方法原型：
// grpc-test-examples/grpctest/server.go
func (s *server) ServerStreamingRPC(req *pb.RequestMessage, stream pb.MyService_ServerStreamingRPCServer) error我們看到除了RequestMessag外，該方法還依賴一個MyService_ServerStreamingRPCServer的類型，這個類型是一個接口類型：
// grpc-test-examples/mygrpc/mygrpc.pb.go
type MyService_ServerStreamingRPCServer interface {
Send(*ResponseMessage) error
grpc.ServerStream
}到這里，你腦子中可能已經冒出了一個想法：使用fake object來對ServerStreamingRPC進行單測[11]，這的確是一個可行的方法，我們下面就基于這個思路實現一下。
注：關于基于fake object進行單測的內容，大家可以看看我以前寫的一篇文章《[]單測時盡量用fake object(https://tonybai.com/2023/04/20/provide-fake-object-for-external-collaborators)》。
3.2 基于fake object的測試
我們首先創建一個實現MyService_ServerStreamingRPCServer的fake object用以代替真實運行RPC服務器時由服務器傳入的stream object：
// grpc-test-examples/grpctest/server_with_fakeobject_test.go
import (
"testing"
pb "demo/mygrpc"
"google.golang.org/grpc"
)
type fakeServerStreamingRPCStream struct {
grpc.ServerStream
responses []*pb.ResponseMessage
}
func (m *fakeServerStreamingRPCStream) Send(resp *pb.ResponseMessage) error {
m.responses = append(m.responses, resp)
return nil
}我們看到fakeServerStreamingRPCStream的Send方法只是將收到的ResponseMessage追加到且內部的ResponseMessage切片中。
接下來我們為ServerStreamingRPC編寫測試用例：
// grpc-test-examples/grpctest/server_with_fakeobject_test.go
func TestServerServerStreamingRPC(t *testing.T) {
s := &server{}
req := &pb.RequestMessage{
Message: "Test message",
}
stream := &fakeServerStreamingRPCStream{}
err := s.ServerStreamingRPC(req, stream)
if err != nil {
t.Fatalf("ServerStreamingRPC failed: %v", err)
}
expectedResponses := []string{
"Server Streaming RPC response 1",
"Server Streaming RPC response 2",
"Server Streaming RPC response 3",
"Server Streaming RPC response 4",
"Server Streaming RPC response 5",
}
if len(stream.responses) != len(expectedResponses) {
t.Errorf("Unexpected number of responses. Got: %d, Want: %d", len(stream.responses), len(expectedResponses))
}
for i, resp := range stream.responses {
if resp.Message != expectedResponses[i] {
t.Errorf("Unexpected response at index %d. Got: %s, Want: %s", i, resp.Message, expectedResponses[i])
}
}
}在這個測試中，ServerStreamingRPC接收一個請求(req)，并通過fake stream object的Send方法返回了5個response，通過與預期的response對比，即可做出測試是否通過的斷言。
到這里，我們看到：fake object完全滿足對gRPC Server Handler進行測試的要求。不過我們需要針對不同的Handler建立不同的fake object類型，和文初基于httptest創建的測試用例相比，用例間欠缺了一些一致性。
那grpc-go是否提供了類似httptest的工具來幫助我們更一致的實現grpc server handler的測試用例呢？我們繼續往下看。
3.3 利用grpc-go提供的測試工具包
grpc-go項目在test下提供了bufconn包，可以幫助我們像httptest那樣建立用于測試的“虛擬gRPC服務器”，下面是基于bufconn包建立gRPC測試用服務器的代碼：
// grpc-test-examples/grpctest/server_with_buffconn_test.go
package main
import (
"context"
"log"
"net"
"testing"
pb "demo/mygrpc"
"google.golang.org/grpc"
"google.golang.org/grpc/test/bufconn"
)
func newGRPCServer(t *testing.T) (pb.MyServiceClient, func()) {
// 創建 bufconn.Listener 作為服務器的監聽器
listener := bufconn.Listen(1024 * 1024)
// 創建 gRPC 服務器
srv := grpc.NewServer()
// 注冊服務處理程序
pb.RegisterMyServiceServer(srv, &server{})
// 在監聽器上啟動服務器
go func() {
if err := srv.Serve(listener); err != nil {
t.Fatalf("Server failed to start: %v", err)
}
}()
// 創建 bufconn.Dialer 作為客戶端連接
dialer := func(context.Context, string) (net.Conn, error) {
return listener.Dial()
}
// 使用 DialContext 和 bufconn.Dialer 創建客戶端連接
conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(dialer), grpc.WithInsecure())
if err != nil {
t.Fatalf("Failed to dial server: %v", err)
}
// 創建客戶端實例
client := pb.NewMyServiceClient(conn)
return client, func() {
err := listener.Close()
if err != nil {
log.Printf("error closing listener: %v", err)
}
srv.Stop()
}
}newGRPCServer是一個用于在測試中創建gRPC服務器和客戶端的輔助函數，它使用bufconn.Listen創建一個bufconn.Listener作為服務器的監聽器。bufconn包提供了一種在內存中模擬網絡連接的方法。然后，它使用grpc.NewServer()創建了一個新的gRPC服務器實例，并使用pb.RegisterMyServiceServer將待測的服務實例(這里是server類型實例)注冊到gRPC服務器中。接下來，它創建了與該服務器建連的gRPC客戶端，由于該客戶端要與bufconn.Listener建連，這里用了一個dialer函數，該函數將通過調用listener.Dial()來建立與服務器的連接。之后基于該連接，我們創建了MyServiceClient的客戶端實例，并返回，供測試用例使用。
基于newGPRCServer這種方式，我們改造一下UnaryRPC的測試用例：
// grpc-test-examples/grpctest/server_with_buffconn_test.go
func TestServerUnaryRPCWithBufConn(t *testing.T) {
client, shutdown := newGRPCServer(t)
defer shutdown()
tests := []struct {
name string
requestMessage *pb.RequestMessage
expectedResp *pb.ResponseMessage
}{
{
name: "Test Case 1",
requestMessage: &pb.RequestMessage{
Message: "Test message",
},
expectedResp: &pb.ResponseMessage{
Message: "Unary RPC response",
},
},
// Add more test cases as needed
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resp, err := client.UnaryRPC(context.Background(), tt.requestMessage)
if err != nil {
t.Fatalf("UnaryRPC failed: %v", err)
}
if resp.Message != tt.expectedResp.Message {
t.Errorf("Unexpected response. Got: %s, Want: %s", resp.Message, tt.expectedResp.Message)
}
})
}
}我們看到，相對于前面的TestServerUnaryRPCs，兩者復雜度在一個層次。如果結合下面的ServerStreamRPC的測試用例，你就能看出這種方式在測試用例一致性方面的優勢了：
// grpc-test-examples/grpctest/server_with_buffconn_test.go
func TestServerServerStreamingRPCWithBufConn(t *testing.T) {
client, shutdown := newGRPCServer(t)
defer shutdown()
req := &pb.RequestMessage{
Message: "Test message",
}
stream, err := client.ServerStreamingRPC(context.Background(), req)
if err != nil {
t.Fatalf("ServerStreamingRPC failed: %v", err)
}
expectedResponses := []string{
"Server Streaming RPC response 1",
"Server Streaming RPC response 2",
"Server Streaming RPC response 3",
"Server Streaming RPC response 4",
"Server Streaming RPC response 5",
}
gotResponses := []string{}
for {
resp, err := stream.Recv()
if err != nil {
break
}
gotResponses = append(gotResponses, resp.Message)
}
if len(gotResponses) != len(expectedResponses) {
t.Errorf("Unexpected number of responses. Got: %d, Want: %d", len(gotResponses), len(expectedResponses))
}
for i, resp := range gotResponses {
if resp != expectedResponses[i] {
t.Errorf("Unexpected response at index %d. Got: %s, Want: %s", i, resp, expectedResponses[i])
}
}
}我們再也無需為每個Server Handler建立各自的fake object了！
由此看到：grpc-go的test/bufconn就是類似httptest的那個grpc server handler的測試腳手架搭建工具。
3.4 其他Streaming模式的Handler測試
有了bufconn這一利器，其他Streaming模式的Handler測試實現邏輯就大同小異了。本文示例中的ClientStreamingRPC和BidirectionalStreamingRPC兩個Handler的測試用例就作為作業，交給各位讀者去完成吧！
4. 小結
在本文中，我們詳細探討了如何對gRPC服務端Handler進行單元測試，我們的目標是找到像net/http/httptest包那樣的，可以為gRPC服務端handler測試提供腳手架代碼幫助的測試方法。
我們按照gRPC的四種通信方式，由簡到難的逐一探討各種Handler的單測方法。UnaryRPC handler測試最為簡單，毫無技巧的普通測試邏輯便能應付。
但一旦涉及streaming通信方式的測試，我們就需要借助類似fake object的單測技術了。但fake object也有不足，那就是需要為每個RPC handler建立單獨的fake object，費時費力還缺少一致性！
好在，grpc-go項目為我們提供了test/bufconn包，該包可以像net/http/httptest包那樣幫助我們快速建立可復用的測試腳手架代碼，這樣我們便可以為所有服務端RPC Handler建立一致、穩定的單元測試用例了！
當然，服務端RPC Handler的單測方法可能不止文中提及這些，各位讀者如果有更好的方法和實踐，歡迎在評論區留言！
本文涉及的源碼可以在這里[12]下載。
5. 參考資料
- Testing gRPC methods[13] - https://medium.com/@johnsiilver/testing-grpc-methods-6a8edad4159d
- 《gRPC Up and Running》[14] - https://book.douban.com/subject/34796013/
- Mocking the Universe: Two Techniques for Testing gRPC with Mocks[15] - https://rotational.io/blog/mocking-the-universe/
