Go: WebSockets 單元測試
WebSockets 通過 TCP 連接提供客戶端與服務器之間的雙向即時通信。這意味着，我們可以維護一個 TCP 連接，然後發送和監聽該連接上的消息，而不是不斷地通過新建 TCP 連接去輪詢 web 服務器的更新。
在 Go 的生態中，WebSocket 協議有幾個不同的實現。有些庫是協議的純實現。另外一些則選擇在 WebSocket 協議的基礎上構建，爲他們特定的用例創建更好的抽象。
下面是一個不全的 Go WebSocket 協議實現列表:
-
x/net/websocket(from Go sub-repository packages) -
gorilla/websocket -
nhooyr/websocket -
gobwas/ws -
sacOO7/gowebsocket
本文我們將使用 gorilla/websocket 庫來實現 websocket 協議。你會發現 websocket 單測與 HTTP 服務的並沒有很大不同。儘管如此，我們在測試時還是要考慮 WebSockets 的一些特殊方面。
WebSocket 在線物品拍賣示例
在線拍賣是以實時通信爲核心的行業之一。在一場拍賣中，幾秒鐘的時間就決定了你是贏了還是失去了一件你一直想要的收藏品。
讓我們以 gorilla/websocket 庫實現的簡單拍賣應用程序作爲本文的示例。
首先，我們將定義兩個非常簡單的結構體 Bid 和 Auction，我們將在 WebSocket 處理程序中使用它們。Auction有一個 Bid 方法，我們將使用該方法接收客戶端發送來的競價請求。
結構體定義
type Bid struct {
UserID int `json:"user_id"`
Amount float64 `json:"amount"`
}
type Auction struct {
ItemID int `json:"item_id"`
EndTime int64 `json:"end_time"`
Bids []*Bid
}
func NewAuction(d time.Duration, itemID int, b []*Bid) Auction {
return Auction{
ItemID: itemID,
EndTime: time.Now().Add(d).Unix(),
Bids: b,
}
}這兩種類型都相當簡單，包含的字段非常少。NewAuction 構造函數創建了一個帶有 duration 拍賣持續時間、itemID 和 * Bids 的 Aution 實例。
競拍
我們將通過Bid方法來實現拍賣的競標動作：
func (a *Auction) Bid(amount float64, userID int) (*Bid, error) {
if len(a.Bids) > 0 {
largestBid := a.Bids[len(a.Bids)-1]
if largestBid.Amount >= amount {
return nil, fmt.Errorf("競拍價必須大於 %.2f", largestBid.Amount)
}
}
if a.EndTime < time.Now().Unix() {
return nil, fmt.Errorf("拍賣已結束")
}
bid := Bid{
Amount: amount,
UserID: userID,
}
// Mutex lock
a.Bids = append(a.Bids, &bid)
// Mutex unlock
return &bid, nil
}Auction 的 Bid 方法就是物品競拍發生的地方。它接收一個amount和userID作爲參數，並向Auction對象中添加 Bid 實例。而且它會檢查競拍是否結束以及客戶端發送的競拍價格是否大於已有的最大競價。如果這些條件中的任何一個不滿足，它將向客戶端返回適當的錯誤。
有了結構體定義和 Bid 方法，讓我們深入到 WebSockets 機制。
WebSocket 連接處理
想象一下，一個可以在拍賣中實時出價的網站。它通過 WebSockets 發送的每一條 JSON 消息都會包含用戶的標識符 (UserID) 和出價的金額 (amount)。一旦服務器接收了消息，它將參與競價並向客戶端返回一個競拍結果。
在服務器端，此通信將由net/http處理程序完成。它將處理所有 WebSocket 的業務邏輯，有幾個值得注意的步驟:
1、將接收到的 HTTP 連接升級爲 WebSocket 連接。
2、接收來自客戶端的消息。
3、從消息中解碼出 bid 對象。
4、參與競價。
5、 向客戶端發送競拍結果。
下面我們來實現這個處理程序。首先定義inbound和outbound消息類型，用於接收和發送客戶端消息。
type inbound struct {
UserID int `json:"user_id"`
Amount float64 `json:"amount"`
}
type outbound struct {
Body string `json:"body"`
}它們都分別表示入站 / 出站消息，這就是在客戶端和服務器之間的交互數據。inbound入站消息將表示一個競標內容，而outbound類型表示一個簡單的返回消息，其 Body 中包含一些文本。
接下來定義bidsHandler，包含 ServeHTTP 方法實現 HTTP 連接的升級：
var upgrader = websocket.Upgrader{}
type bidsHandler struct {
auction *Auction
}
func (bh bidsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
upgrader.CheckOrigin = func(r *http.Request) bool { return true }
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Println("upgrade:", err)
return
}
defer ws.Close()
// 剩餘代碼在後面
}首先定義websocket.Upgrader，接收處理程序的http.ResponseWriter和*http.Resquest並升級連接。因爲這只是一個應用程序示例upgrader.CheckOrigin方法將只返回 true，而不檢查傳入請求的來源。一旦upgrader完成連接的升級，將返回*websocket.Conn對象保存在ws變量中。*websocket.Conn將接收所有客戶端發送來的消息，也是處理程序讀取請求內容的地方。同樣，處理程序將會向*websocket.Conn寫入消息，它將向客戶端發送響應消息。
func (bh bidsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// 前面的代碼...
for {
_, m, err := ws.ReadMessage()
if err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
log.Printf("error: %v", err)
}
return
}
var in inbound
err = json.Unmarshal(m, &in)
if err != nil {
handleError(ws, err)
continue
}
bid, err := bh.auction.Bid(in.Amount, in.UserID)
if err != nil {
handleError(ws, err)
continue
}
out, err := json.Marshal(outbound{Body: fmt.Sprintf("Bid placed: %.2f", bid.Amount)})
if err != nil {
handleError(ws, err)
continue
}
err = ws.WriteMessage(websocket.BinaryMessage, out)
if err != nil {
handleError(ws, err)
continue
}
}
}for循環做了幾件事：首先，使用ws.ReadMessage()讀取 websocket 消息，該函數返回消息類型 (二進制或文本) 和消息內容(m) 以及可能發生的錯誤 (err)。然後，檢查客戶端是否意外地關閉了連接。
錯誤處理完成並讀取到消息，我們將使用json.Unmarshal對其進行解碼。接着調 Bid 方法參與競拍。然後使用json.Marshal對返回內容進行序列化，使用ws.WriteMessage方法發送給客戶端。
測試 WebSockets 處理函數
儘管編寫 WebSocket 處理程序比普通 HTTP 處理程序要複雜得多，但測試它們很簡單。事實上，測試 WebSockets 處理程序就像測試 HTTP 處理程序一樣簡單。這是因爲 WebSockets 是在 HTTP 上構建的，所以測試 WebSockets 使用的工具與測試 HTTP 服務器相同。
首先添加測試用例：
func TestBidsHandler(t *testing.T) {
tcs := []struct {
name string
bids []*Bid
duration time.Duration
message inbound
reply outbound
}{
{
name: "with good bid",
bids: []*Bid{},
duration: time.Hour * 1,
message: inbound{UserID: 1, Amount: 10},
reply: outbound{Body: "Bid placed: 10.00"},
},
{
name: "with bad bid",
bids: []*Bid{
&Bid{
UserID: 1,
Amount: 20,
},
},
duration: time.Hour * 1,
message: inbound{UserID: 1, Amount: 10},
reply: outbound{Body: "amount must be larger than 20.00"},
},
{
name: "good bid on expired auction",
bids: []*Bid{
&Bid{
UserID: 1,
Amount: 20,
},
},
duration: time.Hour * -1,
message: inbound{UserID: 1, Amount: 30},
reply: outbound{Body: "auction already closed"},
},
}
for _, tt := range tcs {
t.Run(tt.name, func(t *testing.T) {
a := NewAuction(tt.duration, 1, tt.bids)
h := bidsHandler{&a}
// 剩餘代碼在後面
})
}
}首先，我們從定義測試用例開始。每個用例有一個name，這是測試用例的可讀名稱。此外，每個測試用例都有一個bids切片和一個 duration 持續時間，用於創建一個測試拍賣對象Auction。測試用例還有一個入站消息inbound和一個出站回覆outbound—這是測試用例將發送給處理程序並期望從處理程序返回的消息。
在 TestBidsHandler 中我們添加三種不同的測試用例——一個是客戶端發起了錯誤的報價，低於目前最大報價，另一個測試用例，客戶端添加了一個正常的報價，第三個客戶端參與的拍賣已結束。
下面完成測試函數：
func TestBidsHandler(t *testing.T) {
// 測試用例和其他內容在前面...
for _, tt := range tcs {
t.Run(tt.name, func(t *testing.T) {
a := NewAuction(tt.duration, 1, tt.bids)
h := bidsHandler{&a}
s, ws := newWSServer(t, h)
defer s.Close()
defer ws.Close()
sendMessage(t, ws, tt.message)
reply := receiveWSMessage(t, ws)
if reply != tt.reply {
t.Fatalf("Expected '%+v', got '%+v'", tt.reply, reply)
}
})
}
}我們在 subtest 函數體中添加了一些新函數。newWSServer將創建一個測試服務器並將其升級爲 WebSocket 連接，同時返回服務器和 WebSocket 連接。然後，sendMessage函數通過 WebSocket 連接將消息從測試用例發送到測試服務器。之後，通過receiveWSMessage，我們將從服務器讀取響應，並通過將其與測試用例的進行比較來斷言其正確性。
那麼，這些新的函數的作用是什麼呢？讓我們逐一分析。
func newWSServer(t *testing.T, h http.Handler) (*httptest.Server, *websocket.Conn) {
t.Helper()
s := httptest.NewServer(h)
wsURL := httpToWs(t, s.URL)
ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatal(err)
}
return s, ws
}newWSServer函數使用httptest.NewServer函數將處理程序掛載到測試 HTTP 服務器上。通過httpToWS，實現了將服務器的URL轉爲 websocket URL (它只是將 URL 中的http協議替換爲ws，或將https替換爲wss協議)。
爲了建立 WebSocket 連接，我們使用WebSocket.DefaultDialer，它是一個所有字段都設置爲默認值的 dialer。調用Dial方法通過 WebSocket 服務器 URL (wsURL) 返回 WebSocket 連接。
func sendMessage(t *testing.T, ws *websocket.Conn, msg inbound) {
t.Helper()
m, err := json.Marshal(msg)
if err != nil {
t.Fatal(err)
}
if err := ws.WriteMessage(websocket.BinaryMessage, m); err != nil {
t.Fatalf("%v", err)
}
}sendMessage函數接收一個 WebSocket 連接和inbound消息作爲參數。將消息序列化成 json 以二進制格式在 websocket 連接中發送。
func receiveWSMessage(t *testing.T, ws *websocket.Conn) outbound {
t.Helper()
_, m, err := ws.ReadMessage()
if err != nil {
t.Fatalf("%v", err)
}
var reply outbound
err = json.Unmarshal(m, &reply)
if err != nil {
t.Fatal(err)
}
return reply
}receiveWSMessage函數以ws WebSocket 連接爲參數，通過ws.ReadMessage()讀取請求消息，然後反序列化成outbound類型返回。
如果我們運行測試，我們將看到它們通過:
$ go test ./... -v
=== RUN TestBidsHandler
=== RUN TestBidsHandler/with_good_bid
=== RUN TestBidsHandler/with_bad_bid
=== RUN TestBidsHandler/good_bid_on_expired_auction
--- PASS: TestBidsHandler (0.00s)
--- PASS: TestBidsHandler/with_good_bid (0.00s)
--- PASS: TestBidsHandler/with_bad_bid (0.00s)
--- PASS: TestBidsHandler/good_bid_on_expired_auction (0.00s)
PASS
ok github.com/fteem/go-playground/testing-in-go-web-sockets 0.013s本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/Jiqoj3ZdsSayC3Y1gkfbag