一文搞懂Go gRPC服務(wù)Handler單元測試

作者：白明的贊賞賬戶 2023-11-25 09:41:34

在本文中，我們詳細(xì)探討了如何對gRPC服務(wù)端Handler進(jìn)行單元測試，我們的目標(biāo)是找到像net/http/httptest包那樣的，可以為gRPC服務(wù)端handler測試提供腳手架代碼幫助的測試方法。

在云原生時(shí)代和微服務(wù)架構(gòu)背景下，HTTP和RPC協(xié)議成為服務(wù)間通信和與客戶端交互的兩種主要方式。對于Go語言而言，標(biāo)準(zhǔn)庫提供了net/http/httptest包，為開發(fā)人員提供了便捷的方式來構(gòu)建服務(wù)端HTTP Handler單元測試的測試腳手架代碼，而無需真正建立HTTP服務(wù)器，讓開發(fā)人員可以聚焦于對Handler業(yè)務(wù)邏輯的測試。比如下面這個(gè)示例：

// grpc-test-examples/httptest/http_handler_test.go

func myHandler(w http.ResponseWriter, r *http.Request) {
    // 設(shè)置響應(yīng)頭
    w.Header().Set("Content-Type", "text/plain")

    // 根據(jù)請求方法進(jìn)行不同的處理
    switch r.Method {
    case http.MethodGet:
        // 處理GET請求
        fmt.Fprint(w, "Hello, World!")
    ... ...
    }
}

func TestMyHandler(t *testing.T) {
    // 創(chuàng)建一個(gè)ResponseRecorder來記錄Handler的響應(yīng)
    rr := httptest.NewRecorder()

    // 創(chuàng)建一個(gè)模擬的HTTP請求，可以指定請求的方法、路徑、正文等
    req, err := http.NewRequest("GET", "/path", nil)
    if err != nil {
        t.Fatal(err)
    }

    // 調(diào)用被測試的Handler函數(shù)，傳入ResponseRecorder和Request對象
    // 這里假設(shè)被測試的Handler函數(shù)為myHandler
    myHandler(rr, req)

    // 檢查響應(yīng)狀態(tài)碼和內(nèi)容
    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來構(gòu)建一個(gè)fake的http server[1]。

然而，對于使用主流的gRPC等RPC協(xié)議的服務(wù)端Handler[2]來說，是否存在類似httptest的測試腳手架生成工具包呢？對gRPC的服務(wù)端Handler有哪些單元測試的方法呢？在這篇文章中，我們就一起來探究一下。

1. 建立被測的gRPC服務(wù)端Handler

我們首先來建立一個(gè)涵蓋多種gRPC通信模式的服務(wù)端Handler集合。

gRPC支持四種通信模式，它們分別為：

簡單RPC(Simple RPC，也稱為Unary RPC)

這是最簡單的，也是最常用的gRPC通信模式，簡單來說就是一請求一應(yīng)答。

服務(wù)端流RPC(Server-streaming RPC)

客戶端發(fā)來一個(gè)請求，服務(wù)端通過流返回多個(gè)應(yīng)答。

客戶端流RPC(Client-streaming RPC)

客戶端通過流發(fā)來多個(gè)請求，服務(wù)端以一個(gè)應(yīng)答回復(fù)。

雙向流RPC(Bidirectional-Streaming RPC)

客戶端通過流發(fā)起多個(gè)請求，服務(wù)端也通過流對應(yīng)返回多個(gè)應(yīng)答。

注：關(guān)于gRPC四種通信方式的詳情，可以參考我之前寫的《gRPC客戶端的那些事兒[3]》一文。

我們這個(gè)SUT(被測目標(biāo))是包含以上四種通信模式的gRPC服務(wù)，它的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

注：你的環(huán)境下需要安裝protoc[5]和protoc-gen-go[6]才能正確執(zhí)行上面生成命令，具體的安裝方法可參考protoc安裝文檔[7]。
注：除了使用經(jīng)典的protoc[8]基于proto文件生成Go源碼外，也可以基于Go開發(fā)的buf cli[9]進(jìn)行代碼生成和API管理。buf cLi是現(xiàn)代、快速、高效的Protobuf API管理的終極工具，為基于Protobuf的開發(fā)和維護(hù)提供了全面的解決方案。等有機(jī)會的時(shí)候，我在以后的文章中詳細(xì)說說buf。

有了生成的樁代碼后，我們便可以建立一個(gè)gRPC服務(wù)器：

// grpc-test-examples/grpctest/main.go

package main
  
import (
    pb "demo/mygrpc"
    "log"
    "net"

    "google.golang.org/grpc"
)

func main() {
    // 創(chuàng)建 gRPC 服務(wù)器
    lis, err := net.Listen("tcp", ":50051")
    if err != nil {
        log.Fatalf("failed to listen: %v", err)
    }
    s := grpc.NewServer()

    // 注冊 MyService 服務(wù)
    pb.RegisterMyServiceServer(s, &server{})

    // 啟動 gRPC 服務(wù)器
    log.Println("Starting gRPC server...")
    if err := s.Serve(lis); err != nil {
        log.Fatalf("failed to serve: %v", err)
    }
}

我們看到：在main函數(shù)中，我們創(chuàng)建了一個(gè)TCP監(jiān)聽器，并使用grpc.NewServer()創(chuàng)建了一個(gè)gRPC服務(wù)器。然后，我們通過調(diào)用pb.RegisterMyServiceServer()將server類型的實(shí)例注冊到gRPC服務(wù)器上，以處理來自客戶端的請求。最后，我們啟動gRPC服務(wù)器并監(jiān)聽指定的端口。

上面代碼中注冊到服務(wù)器中的server類型就是實(shí)現(xiàn)了MyService服務(wù)接口的具體類型，它實(shí)現(xiàn)了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
  }
 }
}

在上面代碼中，我們創(chuàng)建了一個(gè)server結(jié)構(gòu)體類型，并實(shí)現(xiàn)了MyService的所有RPC方法。每個(gè)方法都接收相應(yīng)的請求消息，并返回對應(yīng)的響應(yīng)消息。我們的目標(biāo)僅是演示如何對上述gRPC Handler進(jìn)行單元測試，所以這里的實(shí)現(xiàn)邏輯非常簡單。

接下來，我們就來逐一對這些gRPC的Handler方法進(jìn)行單測，我們先從簡單的UnaryRPC方法開始。

2. Unary RPC Handler的單元測試

Unary RPC是最簡單，也是最容易理解的RPC通信模式，即客戶端與服務(wù)端采用一請求一應(yīng)答的模式。server類型的UnaryRPC Handler方法的原型如下：

// grpc-test-examples/grpctest/server.go

func (s *server) UnaryRPC(ctx context.Context, req *pb.RequestMessage) (*pb.ResponseMessage, error)

就像文章開頭做的那個(gè)httpserver的handler單測一樣，我們肯定不想真實(shí)啟動一個(gè)gRPC server，也不想測試gRPC服務(wù)器本身。我們只想測試服務(wù)端handler方法的邏輯是否正確。

觀察一下這個(gè)方法原型，我們發(fā)現(xiàn)它僅依賴兩個(gè)消息結(jié)構(gòu)：RequestMessage和ResponseMessage，這兩個(gè)消息結(jié)構(gòu)是上面基于proto文件自動生成的，這樣我們就可以不借助任何工具包實(shí)現(xiàn)對UnaryRPC handler方法的單測，也無需啟動真實(shí)的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]和表驅(qū)動的測試也非常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這樣簡單那就好了，但實(shí)際上...，好吧，我們繼續(xù)向下看就好了。

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外，該方法還依賴一個(gè)MyService_ServerStreamingRPCServer的類型，這個(gè)類型是一個(gè)接口類型：

// grpc-test-examples/mygrpc/mygrpc.pb.go

type MyService_ServerStreamingRPCServer interface {
    Send(*ResponseMessage) error
    grpc.ServerStream
}

到這里，你腦子中可能已經(jīng)冒出了一個(gè)想法：使用fake object來對ServerStreamingRPC進(jìn)行單測[11]，這的確是一個(gè)可行的方法，我們下面就基于這個(gè)思路實(shí)現(xiàn)一下。

注：關(guān)于基于fake object進(jìn)行單測的內(nèi)容，大家可以看看我以前寫的一篇文章《[]單測時(shí)盡量用fake object(https://tonybai.com/2023/04/20/provide-fake-object-for-external-collaborators)》。

3.2 基于fake object的測試

我們首先創(chuàng)建一個(gè)實(shí)現(xiàn)MyService_ServerStreamingRPCServer的fake object用以代替真實(shí)運(yùn)行RPC服務(wù)器時(shí)由服務(wù)器傳入的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追加到且內(nèi)部的ResponseMessage切片中。

接下來我們?yōu)镾erverStreamingRPC編寫測試用例：

// 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])         
        }                                                                                                         
    }  
}

在這個(gè)測試中，ServerStreamingRPC接收一個(gè)請求(req)，并通過fake stream object的Send方法返回了5個(gè)response，通過與預(yù)期的response對比，即可做出測試是否通過的斷言。

到這里，我們看到：fake object完全滿足對gRPC Server Handler進(jìn)行測試的要求。不過我們需要針對不同的Handler建立不同的fake object類型，和文初基于httptest創(chuàng)建的測試用例相比，用例間欠缺了一些一致性。

那grpc-go是否提供了類似httptest的工具來幫助我們更一致的實(shí)現(xiàn)grpc server handler的測試用例呢？我們繼續(xù)往下看。

3.3 利用grpc-go提供的測試工具包

grpc-go項(xiàng)目在test下提供了bufconn包，可以幫助我們像httptest那樣建立用于測試的“虛擬gRPC服務(wù)器”，下面是基于bufconn包建立gRPC測試用服務(wù)器的代碼：

// 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()) {
 // 創(chuàng)建 bufconn.Listener 作為服務(wù)器的監(jiān)聽器
 listener := bufconn.Listen(1024 * 1024)

 // 創(chuàng)建 gRPC 服務(wù)器
 srv := grpc.NewServer()

 // 注冊服務(wù)處理程序
 pb.RegisterMyServiceServer(srv, &server{})

 // 在監(jiān)聽器上啟動服務(wù)器
 go func() {
  if err := srv.Serve(listener); err != nil {
   t.Fatalf("Server failed to start: %v", err)
  }
 }()

 // 創(chuàng)建 bufconn.Dialer 作為客戶端連接
 dialer := func(context.Context, string) (net.Conn, error) {
  return listener.Dial()
 }

 // 使用 DialContext 和 bufconn.Dialer 創(chuàng)建客戶端連接
 conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(dialer), grpc.WithInsecure())
 if err != nil {
  t.Fatalf("Failed to dial server: %v", err)
 }

 // 創(chuàng)建客戶端實(shí)例
 client := pb.NewMyServiceClient(conn)
 return client, func() {
  err := listener.Close()
  if err != nil {
   log.Printf("error closing listener: %v", err)
  }
  srv.Stop()
 }
}

newGRPCServer是一個(gè)用于在測試中創(chuàng)建gRPC服務(wù)器和客戶端的輔助函數(shù)，它使用bufconn.Listen創(chuàng)建一個(gè)bufconn.Listener作為服務(wù)器的監(jiān)聽器。bufconn包提供了一種在內(nèi)存中模擬網(wǎng)絡(luò)連接的方法。然后，它使用grpc.NewServer()創(chuàng)建了一個(gè)新的gRPC服務(wù)器實(shí)例，并使用pb.RegisterMyServiceServer將待測的服務(wù)實(shí)例(這里是server類型實(shí)例)注冊到gRPC服務(wù)器中。接下來，它創(chuàng)建了與該服務(wù)器建連的gRPC客戶端，由于該客戶端要與bufconn.Listener建連，這里用了一個(gè)dialer函數(shù)，該函數(shù)將通過調(diào)用listener.Dial()來建立與服務(wù)器的連接。之后基于該連接，我們創(chuàng)建了MyServiceClient的客戶端實(shí)例，并返回，供測試用例使用。

基于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，兩者復(fù)雜度在一個(gè)層次。如果結(jié)合下面的ServerStreamRPC的測試用例，你就能看出這種方式在測試用例一致性方面的優(yōu)勢了：

// 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])
  }
 }
}

我們再也無需為每個(gè)Server Handler建立各自的fake object了！

由此看到：grpc-go的test/bufconn就是類似httptest的那個(gè)grpc server handler的測試腳手架搭建工具。

3.4 其他Streaming模式的Handler測試

有了bufconn這一利器，其他Streaming模式的Handler測試實(shí)現(xiàn)邏輯就大同小異了。本文示例中的ClientStreamingRPC和BidirectionalStreamingRPC兩個(gè)Handler的測試用例就作為作業(yè)，交給各位讀者去完成吧！

4. 小結(jié)

我們按照gRPC的四種通信方式，由簡到難的逐一探討各種Handler的單測方法。UnaryRPC handler測試最為簡單，毫無技巧的普通測試邏輯便能應(yīng)付。

但一旦涉及streaming通信方式的測試，我們就需要借助類似fake object的單測技術(shù)了。但fake object也有不足，那就是需要為每個(gè)RPC handler建立單獨(dú)的fake object，費(fèi)時(shí)費(fèi)力還缺少一致性！

好在，grpc-go項(xiàng)目為我們提供了test/bufconn包，該包可以像net/http/httptest包那樣幫助我們快速建立可復(fù)用的測試腳手架代碼，這樣我們便可以為所有服務(wù)端RPC Handler建立一致、穩(wěn)定的單元測試用例了！

當(dāng)然，服務(wù)端RPC Handler的單測方法可能不止文中提及這些，各位讀者如果有更好的方法和實(shí)踐，歡迎在評論區(qū)留言！

本文涉及的源碼可以在這里[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/

責(zé)任編輯：武曉燕來源： TonyBai