Testing A Function That Calls Go Routine

#go #testing

Okay so I've been trying to test a function that calls a routine and I found no solution that fits my taste. This is fairly simple, although some won't like it. Imagine you have a function like this.

package main

import "fmt"

type printer struct{}

func (p printer) HelloWorld() {
  fmt.Println("Hello, world!")
}

type asyncPrinter struct{
  p printer
}

func (ap asyncPrinter) HelloWorld() {
  go ap.p.HelloWorld()
}

Now, if you test this the 'usual' way, it won't wait for the go routine. For example:

package main

import "github.com/stretchr/testify/mock"

type printerMock struct {
  mock.Mock
}

func (pm printerMock) HelloWorld() {
  pm.Called()
}

func Test_AsyncHelloWorld(t *testing.T) {
  pm := printerMock{}
  ap := asyncPrinter(pm)

  pm.On("HelloWorld")
  ap.HelloWorld()

  pm.AssertExpectations(t)
}

This will sometimes fail, the other time will pass. To 'trick' it some will suggest waiting in test. For examples:

package main

import (
  "time"

  "github.com/stretchr/testify/mock"
)

type printerMock struct {
  mock.Mock
}

func (pm printerMock) HelloWorld() {
  pm.Called()
}

func Test_AsyncHelloWorld(t *testing.T) {
  pm := printerMock{}
  ap := asyncPrinter(pm)

  pm.On("HelloWorld")
  ap.HelloWorld()

  time.Sleep(1 * time.Second)
  pm.AssertExpectations(t)
}

Yes, it works! But we don't know how long we should wait till the go routine finish. The way that I like to do it is to use wait group and help of test variable. Take a look at the code:

package main

import (
  "fmt"
  "sync"
)

var testMode bool
var wg sync.WaitGroup

type printer struct{}

func (p printer) HelloWorld() {
  fmt.Println("Hello, world!")
}

type asyncPrinter struct{
  p printer
}

func (ap asyncPrinter) HelloWorld() {
  if testMode {
    wg.Add(1)
  }
  go ap.p.HelloWorld()
  if testMode {
    wg.Wait()
  }
}

This way, we can now simply write test that calls wg.Done() on method call. For examples:

package main

import (
  "sync"

  "github.com/stretchr/testify/mock"
)

type printerMock struct {
  mock.Mock
}

func (pm printerMock) HelloWorld() {
}

func Test_AsyncHelloWorld(t *testing.T) {
  pm := printerMock{}
  ap := asyncPrinter(pm)

  testMode = true
  pm.On("HelloWorld").Run(func(){
    wg.Done()
  })
  ap.HelloWorld()

  pm.AssertExpectations(t)
}