はじめに
goroutine のユニットテストで、期待した関数が呼ばれる前にテストが終了する問題に遭遇しました。
今回は、その対策を紹介します。Go のユニットテストやモックについては、こちらの記事がわかりやすいです。
テスト対象コードの作成
goroutine のテストを書くために簡易的なコードを作成しました。
無限ループで、channel からメッセージを取得して、ログに出力するコードです。
main では、handler.HandleMessage() を実行します。
package main
import (
"github.com/articles/goroutine_test/handler"
)
func main() {
messageHandler := handler.NewHandler()
var forever chan string
go messageHandler.HandleMessage()
<-forever
}
handler では、メッセージの受信を開始し、メッセージ受信時に Info ログを出力します。
メッセージの受信は、for message := range h.messageCh で、行っています。これを、今回のテスト対象のコードとします。
package handler
import (
"log/slog"
"os"
"github.com/articles/goroutine_test/logger"
"github.com/articles/goroutine_test/reciever"
)
type handler struct {
logger logger.ILogger
reciever reciever.IReciever
messageCh chan string
}
func NewHandler() *handler {
return &handler{
logger: slog.New(slog.NewJSONHandler(os.Stdout, nil)),
reciever: reciever.NewReciever(),
messageCh: make(chan string, 1),
}
}
func (h *handler) HandleMessage() {
h.logger.Info("Start Recieve Message")
// メッセージの受信を開始する
go h.reciever.Listen(h.messageCh)
// メッセージを受信したら、ログを出力する
for message := range h.messageCh {
h.logger.Info(message)
}
}
また、handler で使っている logger は、モック化するために、interfaceで定義しておきます。
package logger
type ILogger interface {
Info(string, ...any)
}
reciever は、10 秒ごとに messageCh に "a" を送信します。
package reciever
import "time"
type IReciever interface {
Listen(messageCh chan string)
}
type reciever struct{}
func NewReciever() IReciever {
return &reciever{}
}
func (r *reciever) Listen(messageCh chan string) {
// 10秒ごとにメッセージを受信する
for {
time.Sleep(10 * time.Second)
messageCh <- "a"
}
}
モックの作成
今回は、handler.go のテストを書きます。そこで、handler.go で使用している logger と reciever をモックします。
$ mockgen -source=logger.go -package logger -destination=mock_logger.go
$ mockgen -source=reciever.go -package recieve -destination=mock_reciever.go
失敗するテストコードの作成
以下、テストコードです。
テストの実行ループで、tt.prepareMock() を呼び、テストケースごとにモックの準備をします。
また、tt.prepareChannel(h) を呼び、テストケースに合わせて channel に値を送信します。
package handler
import (
"testing"
"github.com/articles/goroutine_test/logger"
"github.com/articles/goroutine_test/reciever"
"github.com/golang/mock/gomock"
)
func Test_handler_HandleMessage(t *testing.T) {
type fields struct {
logger *logger.MockILogger
reciever *reciever.MockIReciever
}
ctrl := gomock.NewController(t)
mockLogger := logger.NewMockILogger(ctrl)
mockReciever := reciever.NewMockIReciever(ctrl)
tests := []struct {
name string
fields fields
prepareMock func()
prepareChannel func(*handler)
wantErr bool
}{
// テストケース
{
name: "Infoログが出力される",
fields: fields{
logger: mockLogger,
reciever: mockReciever,
},
prepareMock: func() {
mockLogger.EXPECT().Info("Start Recieve Message")
mockReciever.EXPECT().Listen(gomock.Any())
mockLogger.EXPECT().Info("a")
},
prepareChannel: func(handler *handler) {
handler.messageCh <- "a"
},
wantErr: false,
},
}
// テストの実行
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h := &handler{
logger: tt.fields.logger,
reciever: tt.fields.reciever,
messageCh: make(chan string, 1),
}
tt.prepareMock()
tt.prepareChannel(h)
go h.HandleMessage()
})
}
}
テスト実行
上で作成したテストを実行すると、失敗しました。
メッセージにから、gomock の EXPECT() で指定した関数が呼ばれていないことが読み取れます。
これは、go h.HandleMessage() で goroutine を実行した後、期待された関数が呼ばれる前テストが終了してしまうためです。
=== RUN Test_handler_HandleMessage
=== RUN Test_handler_HandleMessage/Infoログが出力される
--- PASS: Test_handler_HandleMessage/Infoログが出力される (0.00s)
/handler/controller.go:137: missing call(s) to *logger.MockILogger.Info(is equal to Start Recieve Message (string)) /handler/handler_test.go:33
/handler/controller.go:137: missing call(s) to *logger.MockILogger.Info(is equal to a (string)) /handler/handler_test.go:35
/handler/controller.go:137: missing call(s) to *reciever.MockIReciever.Listen(is anything) /handler/handler_test.go:34
/handler/controller.go:137: aborting test due to missing call(s)
--- FAIL: Test_handler_HandleMessage (0.00s)
FAIL
FAIL /handler 0.295s
対策
まず、一定時間 Sleep させる方法があると思います。しかし、どの程度待てば良いか明確ではなく、大量のテストケースがある場合に、テストの実行時間が長くなるため、得策ではなさそうです。
今回紹介するのは、モックの中で、sync.WaitGroup.Wait() を使います。以下の変更
を行いました。
- モックの構造体に、
sync.WaitGroupとcalledCountを追加 -
EXPECTされる関数の中で、WaitGroup のカウンタをインクリメント - WaitGroup のカウンタを更新する
SetWaitDelta(int)を追加 - WaitGroup のカウンタが 0 になった後に
calledCountの値を返すGetCalledCount()を追加
Logger のモックにも同様の追加をしていますが、ここでは省略します。
// MockIReciever is a mock of IReciever interface.
type MockIReciever struct {
ctrl *gomock.Controller
recorder *MockIRecieverMockRecorder
+ wg sync.WaitGroup
+ calledCount int
}
// Listen mocks base method.
func (m *MockIReciever) Listen(messageCh chan string) {
+ m.calledCount += 1
m.ctrl.T.Helper()
m.ctrl.Call(m, "Listen", messageCh)
+ m.wg.Done()
}
+func (m *MockIReciever) SetWaitDelta(delta int) {
+ m.wg.Add(delta)
+}
+
+// WaitGroup のカウンタが 0 になるまで待機してから calledCount を返す
+func (m *MockIReciever) GetCalledCount() int {
+ m.wg.Wait()
+ return m.calledCount
+}
次にテストコードを更新します。
prepareMock() の内部で、先ほど追加した SetWaitDelta() を実行し、それぞれのモックの期待実行回数を指定します。そして、テストの実行ループで、mock の calledCount が期待通りの回数であることを確認しています。この時、GetCalledCount() を呼ぶことで、モックで指定した WaitGroup.Wait() が実行され、関数が期待通り呼ばれるまで待機することができます。
import (
"github.com/articles/goroutine_test/logger"
"github.com/articles/goroutine_test/reciever"
"github.com/golang/mock/gomock"
+ "github.com/stretchr/testify/assert"
)
/* -------------------- 省略(変更なし) -------------------- */
tests := []struct {
name string
fields fields
prepareMock func()
prepareChannel func(*handler)
+ expectedRecieverCalledCount int
+ expectedLoggerCalledCount int
wantErr bool
}{
// テストケース
{
name: "Infoログが出力される",
fields: fields{
logger: mockLogger,
reciever: mockReciever,
},
prepareMock: func() {
+ mockReciever.SetWaitDelta(1)
+ mockLogger.SetWaitDelta(2)
mockLogger.EXPECT().Info("Start Recieve Message")
mockReciever.EXPECT().Listen(gomock.Any())
mockLogger.EXPECT().Info("a")
},
prepareChannel: func(handler *handler) {
handler.messageCh <- "a"
},
+ expectedRecieverCalledCount: 1,
+ expectedLoggerCalledCount: 2,
+ wantErr: false,
}
// テストの実行
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h := &handler{
logger: tt.fields.logger,
reciever: tt.fields.reciever,
messageCh: make(chan string, 1),
}
tt.prepareMock()
tt.prepareChannel(h)
go h.HandleMessage()
+ assert.Equal(t, tt.expectedRecieverCalledCount, mockReciever.GetCalledCount())
+ assert.Equal(t, tt.expectedLoggerCalledCount, mockLogger.GetCalledCount())
})
}
/* -------------------- 省略(変更なし) -------------------- */
モックを更新していますが、 mockgen するたびにファイルが上書きされてしまうので、別名で保存した方が良いかも
テスト実行
修正したテストを実行すると、無事テストが pass するようになりました。
=== RUN Test_handler_HandleMessage
=== RUN Test_handler_HandleMessage/Infoログが出力される
--- PASS: Test_handler_HandleMessage/Infoログが出力される (0.00s)
--- PASS: Test_handler_HandleMessage (0.00s)
PASS
ok /handler 0.110s
まとめ
goroutine のテストで関数が呼ばれる前にテストが終了する問題の対策を紹介しました。
本記事では、sync.WaitGroup を活用することで、テストが早期に終了することを防ぐことができました。
最後までお読みいただきありがとうございました。