Contents

How to mock Google BigQuery client in Golang

Whats the problem?

The official Google BigQuery packge for Golang is cloud.google.com/go/bigquery and when you create a client it returns a specific struct instead of an interface. This is a bit confusing, because even though it allow maintainers to add new methods to the bigquery client without worry that it would break consumers, it makes unit testing a bit more complex.

With a specific struct returned you can no longer simply sketch an object that fullfil the interface requirements or use mockery to do it for you. Suppose you have this toy example which prints the data from a table:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
type Row struct {
	value string
	count int64
}

type Aggregator struct {
	bqClient *bigquery.Client
}

func (a *Aggregator) Print(ctx context.Context) error {
	query := a.bqClient.Query("SELECT * FROM `mytestproject-479821.test_dataset.test_table`")

	it, err := query.Read(ctx)
	if err != nil {
		panic(err)
	}

	for {
		var row Row
		err := it.Next(&row)
		if err == iterator.Done { break }
		if err != nil {
			panic(err)
		}
		fmt.Println(row)
	}

	return nil
}

func (a *Aggregator) Close() error {
	return a.bqClient.Close()
}

func NewAggregator(bqClient *bigquery.Client) (*Aggregator, error) {
	return &Aggregator{bqClient}, nil
}

There is no simple way to unit test this code by replacing interation with the BigQuery.

Solution

If you want to test that code you would need to implement an interface for the BigQuery. It is a bit annoying, but it provides you a simple way of replacing your mocks:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
// BigQueryFacade is a simple proxy interface for bigquery client to implement mocks
type BigQueryFacade interface {
    // ReadQuery executes Read+Query calls
	ReadQuery(ctx context.Context, q string) (*BigQueryRowIterator, error)
}

// BigQueryRowIterator is a proxy for bigquery RowIterator
type BigQueryRowIterator interface {
	Next(dst interface{}) error
}

type bigQueryRowIterator struct {
	it *bigquery.RowIterator
}

type bigQueryFacade struct {
	bqClient *bigquery.Client
}

func (a *bigQueryFacade) ReadQuery(ctx context.Context, q string) (BigQueryRowIterator, error) {
	query := a.bqClient.Query(q)

	it, err := query.Read(ctx)
	return &bigQueryRowIterator{
		it: it,
	}, err
}

func (a *bigQueryRowIterator) Next(dst interface{}) error {
	return a.it.Next(dst)
}

// NewBigQueryFacade creates a new proxy object
func NewBigQueryFacade(bqClient *bigquery.Client) *bigQueryFacade {
	return &bigQueryFacade{
		bqClient: bqClient,
	}
}

Now we can create a mocked facade object and use it to unit test our aggregator class

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
type mockBigQuery struct {
	rows    []Row
	queries []string
}

func (m *mockBigQuery) ReadQuery(ctx context.Context, q string) (bq.BigQueryRowIterator, error) {
	m.queries = append(m.queries, q)
	return &mockIterator{
		rows: append([]Row(nil), m.rows...),
	}, nil
}

func (m *mockBigQuery) Close() error {
	return nil
}

type mockIterator struct {
	rows []Row
	idx  int
}

func (m *mockIterator) Next(dst interface{}) error {
	if m.idx >= len(m.rows) {
		return iterator.Done
	}

	row := m.rows[m.idx]
	m.idx++

	switch v := dst.(type) {
	case *Row:
		*v = row
		return nil
	default:
		return fmt.Errorf("unexpected dst type %T", dst)
	}
}

And use it like this in the test cases:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

func TestAggregatorPrint(t *testing.T) {
	ctx := context.Background()
	mockBQ := &mockBigQuery{
		rows: []Row{
			{Value: "foo", Count: 1},
			{Value: "bar", Count: 2},
		},
	}

	a, err := NewAggregator(mockBQ)
	...
}