GoでWebSocketのテスト書く

この記事は、2023年traP新歓ブログリレー5日目の記事です。

こんにちは。20Bの@Rasです。traPでは主にSysAd班で活動しています。
先日traQでのメッセージ数が1位になりました。

pasted-image-20230313220701

東工大に合格された方、おめでとうございます。過去に自炊や一人暮らしのブログを書いているので是非よかったら見てください。
→ 過去のブログ一覧

1年半ほど前にVSCodeで効率的にGoのテストを書く記事を書きました。
→ VSCodeで手を抜いてGoのテストを手を抜かずに書く

ごくたまに反響をもらえるのでありがたい限りです。

今回はGoでWebsocketのテストを書く手法を紹介します。

手元では再現できるんですよね

サンプルを書きました。
https://github.com/ras0q/go-wstest

このレポジトリには2つの実行コマンドがあります。

app: 簡易的なWebSocketサーバーを立ち上げることができるコマンド
- 接続後、クライアントからメッセージを受け取りやまびこを返す実装
- 主な実装はws.Serveに書かれている
wscat: WebSocketサーバーに接続して対話形式で通信を行うコマンド
- JavaScript製のwebsockets/wscatを参考に作っている

$ git clone git@github.com:ras0q/go-wstest.git

$ make
go build -o ./bin/app .
go build -o ./bin/wscat ./wscat

$ tree ./bin
./bin
├── app
└── wscat

0 directories, 2 files

appを実行し localhost:8080 にWebSocketサーバーを立ち上げた後、別のシェルでwscatを実行すると無事やまびこが返ってくるのが確認できます。

$ ./bin/wscat ws://localhost:8080/ws
Connected to ws://localhost:8080/ws
> hello
< You said: hello
> yo
< You said: yo
> aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
< You said: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa

手元でやまびこが返ってくるのは確認できましたが、これをテストするためには少し特殊な方法を使う必要があります。

WebSocketのテスト...？

結論を先に貼っておきます。
https://github.com/ras0q/go-wstest/blob/main/ws/ws_test.go

GoでWebアプリケーションを実装した時、HTTPハンドラーのテストを書くときには、標準パッケージの net/http/httptest を使うのが一般的でしょう。httptest.NewServer(handler http.Handler) *Serverを使うとgoroutineを使っていい感じに別スレッドでHTTPサーバーを立ててくれます。

少し荒業ですが、WebSocketのテストでもこれを使います。

httptest.NewServerの引数にはhttp.Handler インターフェイスを渡す必要があるので、*wsHandlerがhttp.Handlerを満たすように実装します。

type wsHandler struct{}

// *wsHandlerはhttp.Handlerインターフェイスを満たす
func (h *wsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	ws.Serve(w, r)
}

この*wsHandlerを使ってhttptest.Serverを立ち上げます。ついでに自動で作られるURLのスキーマもhttp://からws://に変えてしまいます。

server := httptest.NewServer(&wsHandler{})
server.URL = strings.Replace(server.URL, "http", "ws", 1)

あとはこのURLに向けてWebSocketを接続すると1つのソースコード上でWebSocketのサーバーとクライアントを実装することができます。

c, _, err := websocket.DefaultDialer.Dial(server.URL, nil)
if err != nil {
	...
}

実際にテストを書いてみると以下のようになります。(ws_test.go と内容は同じです)

func TestServe(t *testing.T) {
	t.Run("n回並行にWebsocketの相互通信を行う", func(t *testing.T) {
		t.Parallel()

		var (
			n  = 10
			wg = sync.WaitGroup{}
		)

		wg.Add(n)
		defer wg.Wait()

		for i := 0; i < n; i++ {
			go func(i int) {
				defer wg.Done()

				// httptestを無理やりws用に使う
				server := httptest.NewServer(&wsHandler{})
				server.URL = strings.Replace(server.URL, "http", "ws", 1)

				// クライアントを立ち上げる
				c, _, err := websocket.DefaultDialer.Dial(server.URL, nil)
				require.NoError(t, err)

				// "hello {i}"と送信すると、、、
				err = c.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("hello %d", i)))
				require.NoError(t, err)

				// "You said: hello {i}"と返ってくる
				_, msg, err := c.ReadMessage()
				require.NoError(t, err)
				require.Equal(t, []byte(fmt.Sprintf("You said: hello %d", i)), msg)
			}(i)
		}
	})
}

複数同時接続を意識してgoroutineを回していますが、問題なくテストは成功します。

$ go test ./ws/... -v -race
=== RUN   TestServe
=== RUN   TestServe/n回並行にWebsocketの相互通信を行う
=== PAUSE TestServe/n回並行にWebsocketの相互通信を行う
=== CONT  TestServe/n回並行にWebsocketの相互通信を行う
client1: 2023/03/13 21:21:15 connected
client2: 2023/03/13 21:21:15 connected
client3: 2023/03/13 21:21:15 connected
client4: 2023/03/13 21:21:15 connected
client5: 2023/03/13 21:21:15 connected
client6: 2023/03/13 21:21:15 connected
client7: 2023/03/13 21:21:15 connected
client8: 2023/03/13 21:21:15 connected
client10: 2023/03/13 21:21:15 connected
client9: 2023/03/13 21:21:15 connected
--- PASS: TestServe (0.00s)
    --- PASS: TestServe/n回並行にWebsocketの相互通信を行う (0.01s)
PASS
ok      github.com/ras0q/go-wstest/ws

終わりに

GoでWebSocketのテストを書いてる記事が見つからなかったので書いたんですが、この記事を書いている途中で見つけてしまいました。共有だけしておきます。

https://ieftimov.com/posts/testing-in-go-websockets/

新入生の中には大学からプログラミングを始めようと思っている方も多いと思います。筆者もその一人でした。最初からこの記事を完全に理解する必要はないです。traPでは初心者のサポートがめちゃめちゃ手厚いので是非体験会だけでもお越しください！

明日は、@kashiwadeさんの記事です。楽しみ〜

GoでWebSocketのテスト書く

手元では再現できるんですよね

WebSocketのテスト...？

終わりに

この記事をシェア

関連する記事