はじめに

goroutine のユニットテストで、期待した関数が呼ばれる前にテストが終了する問題に遭遇しました。
今回は、その対策を紹介します。Go のユニットテストやモックについては、こちらの記事がわかりやすいです。

テスト対象コードの作成

goroutine のテストを書くために簡易的なコードを作成しました。
無限ループで、channel からメッセージを取得して、ログに出力するコードです。

main では、handler.HandleMessage() を実行します。

main.go

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 で、行っています。これを、今回のテスト対象のコードとします。

handler/handler.go

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で定義しておきます。

logger/logger.go

package logger

type ILogger interface {
	Info(string, ...any)
}

reciever は、10 秒ごとに messageCh に "a" を送信します。

reciever/reciever.go

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 に値を送信します。

handler/handler_test.go

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 のモックにも同様の追加をしていますが、ここでは省略します。

reciever/mock_reciever.go

// 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() が実行され、関数が期待通り呼ばれるまで待機することができます。

handler/handler_test.go

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 を活用することで、テストが早期に終了することを防ぐことができました。

最後までお読みいただきありがとうございました。