Unit Testing and mock calls inside goroutines

Recently I came across a situation where I had to expect a call to a method that was being called inside a goroutine. I used gomock to mock interfaces for the unit tests.

The use case was to post a webhook after a successful update in a goroutine.

Let’s get to the details of writing tests by mocking the webhook client.

Following is the code for the webhook client. For the sake of simplicity, I have used fmt.Printf instead of actual network call.

package clients

//go:generate mockgen -source=clients.go -destination=./clients_mock/clients_mock.go -package=clients_mock WebhookInterface


import "fmt"

type WebhookInterface interface {
   PostWebhook(messageID int, description string)
}

type WebhookClient struct {}

func NewWebhookClient() *WebhookClient {
   return &WebhookClient{}
}

func (w *WebhookClient) PostWebhook(messageID int, description string) {
   fmt.Printf("posting webhook messageID: %d, description: %s", messageID, description )
}

We can generate mock for WebhookInterface using mockgen. The following code stub is used to generate the mocks.

//go:generate mockgen -source=clients.go -destination=./clients_mock/clients_mock.go -package=clients_mock WebhookInterface

To generate the mocks across the project, run go generate ./... from the project root directory.

The code snippet below shows the actual use of the WebhookInterface

type UpdateService struct {
   webHookClient clients.WebhookInterface
}

func NewUpdateService(webhookCli clients.WebhookInterface) *UpdateService {
   return &UpdateService{
      webHookClient: webhookCli,
   }
}

func(s *UpdateService) Update(messageID int, description string) bool {
   fmt.Printf("updating with messageID: %d, description: %s", messageID, description)
   go s.processAndPostWebhook(messageID, description)
   return true
}

func (s *UpdateService)processAndPostWebhook(messageID int, description string){
   time.Sleep(100 * time.Second)
   s.webHookClient.PostWebhook(messageID+1, description)
}

I have introduced time.Sleep(100 *time.Millisecond) to have a delay before calling PostWebhook in a goroutine.

This will work without issues with the main program since the main routine started an HTTP server listening on a port. This allows the goroutine to get scheduled and executed as expected. However, this is not the case while running a unit test.

The Problem

We would usually write unit tests for the function Update() in the following manner.

func TestUpdateService_Update(t *testing.T) {
      ctrl := gomock.NewController(t)
      defer ctrl.Finish()
      // create WebhookClient mock
      mockWebhookClient := clients_mock.NewMockWebhookInterface(ctrl)
      messageID := 1
      description := "Updated message"
      // expect a call to PostWebhook
      mockWebhookClient.EXPECT().
      PostWebhook(messageID+1, description)
      updateService := services.NewUpdateService(mockWebhookClient)
      assert.True(t, updateService.Update(messageID, description))
}

The above test fails with the following error, even though expecting a call to PostWebhook was set.

missing call(s) to *clients_mock.MockWebhookInterface.PostWebhook(is equal to 1, is equal to Updated message) services_test.go:19
...
aborting test due to missing call(s)
--- FAIL: TestUpdateService_Update (0.00s)

This is causing because PostWebhook is being called in a goroutine and the test exits before the go routine processAndPostWebhook gets scheduled and could call PostWebhook , hence ctrl.Finish() will fail the test.

The Solution

To make the test pass, we need to make use of the Do method provided by gomock along withsync.WaitGroup

Here’s the working solution, the test that passes and matches the expectations

func TestUpdateService_Update(t *testing.T) {
      ctrl := gomock.NewController(t)
      defer ctrl.Finish()
      // create WebhookClient mock
      mockWebhookClient := clients_mock.NewMockWebhookInterface(ctrl)
      messageID := 1
      description := "Updated message"
      // expect a call to PostWebhook
      wg := sync.WaitGroup{}
      wg.Add(1)
      mockWebhookClient.EXPECT().
      PostWebhook(messageID+1, description).
      Do(func(arg0, arg1 interface{}) {
         defer wg.Done()
      })
      updateService := services.NewUpdateService(mockWebhookClient)
      isUpdated := updateService.Update(messageID, description)
      wg.Wait()
      assert.True(t, isUpdated)
}

Do function is executed after when the call is matched for the expectation. Here we make use of Waitgroup to wait until a call is matched for PostWebhook(messageID+1, description) . The Do function executes only after the expectation for PostWebhook is matched.

The key thing to note here is that Do method expects a function with the same signature as that of the EXPECT -ing function. In our case PostWebhook(arg0, arg1 interface{}

The test will fail if the signature of the function passed to theDo method doesn’t match the signature of PostWebhook(arg0, arg1 interface{}

I hope you find this article helpful! Please leave your thoughts in the responses section.