Unit Testing Golang Code Calling Elasticsearch

There is an amazing Elasticsearch client for Go and in Working With Elasticsearch I demonstrated with example how to index and search for documents using it. However, if you want to ensure your code will continue working correctly and will not be inadvertently broken during refactoring, you have to cover the code with tests.

In this post I’ll show how to unit test Go code interoperating with Elasticsearch. Keep in mind, however, that the same approach is applicable to unit testing almost any code calling external RESTful API.

A service calling Elasticsearch

Suppose you have a logging service which can fetch n last entries of any application log sorted by time in ascending order. Method GetLog below does exactly that. I provided whole code for a production ready service, to show you how it usually looks in practice.

package logging

import (
	"gopkg.in/olivere/elastic.v3"
	"reflect"
)

type Service interface {
	GetLog(app string, lines int) ([]string, error)
}

func NewService(url string) (Service, error) {
	client, err := elastic.NewSimpleClient(elastic.SetURL(url))
	if err != nil {
		return nil, err
	}
	return &service{elasticClient: client}, nil
}

type service struct {
	elasticClient *elastic.Client
}

type Log struct {
	Message string `json:"message"`
}

// GetLog returns limited tail of log sorted by time in ascending order
func (s *service) GetLog(app string, limit int) ([]string, error) {
	termQuery := elastic.NewTermQuery("app", app)

	res, err := s.elasticClient.Search("_all").
		Query(termQuery).
		Sort("@timestamp", false).
		Size(limit).
		Do()

	if err != nil {
		return nil, err
	}

	msgNum := len(res.Hits.Hits)
	if msgNum == 0 {
		return []string{}, nil
	}

	messages := make([]string, msgNum, msgNum)

	var l Log
	for i, item := range res.Each(reflect.TypeOf(l)) {
		l := item.(Log)
		messages[i] = l.Message
	}

	// Reversing messages
	for i := 0; i < msgNum/2; i++ {
		messages[i], messages[msgNum-(i+1)] = messages[msgNum-(i+1)], messages[i]
	}

	return messages, nil
}

Note that log is fetched from Elasticsearch in descending order first and then reversed back before returning result to caller of logging service. This trick is needed to ensure we get tail of the log capped by limit.

Unit testing the service

In general, the simplest way to unit test code calling external service via client is to mock the client. Unfortunately, elastic.Client is implemented as struct, which makes mocking it a not trivial issue.

Much easier in this case would be to go deeper and mock Elasticsearch RESTful API. One of options is to run httptest.Server which will serve a hardcoded response, where the response is a copy of what Elasticsearch fetches for some predefined request:

package logging

import (
	"github.com/stretchr/testify/assert"
	"gopkg.in/olivere/elastic.v3"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestLog(t *testing.T) {
	handler := http.NotFound
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		handler(w, r)
	}))
	defer ts.Close()

	handler = func(w http.ResponseWriter, r *http.Request) {
		resp := `{
			  "took" : 122,
			  "timed_out" : false,
			  "_shards" : {
			    "total" : 6,
			    "successful" : 5,
			    "failed" : 1,
			    "failures" : [ {
			      "shard" : 0,
			      "index" : ".kibana",
			      "node" : "jucBX9QkQIini9dLG9tZIw",
			      "reason" : {
				"type" : "search_parse_exception",
				"reason" : "No mapping found for [offset] in order to sort on"
			      }
			    } ]
			  },
			  "hits" : {
			    "total" : 10,
			    "max_score" : null,
			    "hits" : [ {
			      "_index" : "logstash-2016.07.25",
			      "_type" : "log",
			      "_id" : "AVYkNv542Gim_t2htKPU",
			      "_score" : null,
			      "_source" : {
				"message" : "Alice message 10",
				"@version" : "1",
				"@timestamp" : "2016-07-25T22:39:55.760Z",
				"source" : "/Users/yury/logs/alice.log",
				"offset" : 144,
				"type" : "log",
				"input_type" : "log",
				"count" : 1,
				"fields" : null,
				"beat" : {
				  "hostname" : "Yurys-MacBook-Pro.local",
				  "name" : "Yurys-MacBook-Pro.local"
				},
				"host" : "Yurys-MacBook-Pro.local",
				"tags" : [ "beats_input_codec_plain_applied" ],
				"app" : "alice"
			      },
			      "sort" : [ 144 ]
			    }, {
			      "_index" : "logstash-2016.07.25",
			      "_type" : "log",
			      "_id" : "AVYkNv542Gim_t2htKPT",
			      "_score" : null,
			      "_source" : {
				"message" : "Alice message 9",
				"@version" : "1",
				"@timestamp" : "2016-07-25T22:39:55.760Z",
				"source" : "/Users/yury/logs/alice.log",
				"offset" : 128,
				"input_type" : "log",
				"count" : 1,
				"beat" : {
				  "hostname" : "Yurys-MacBook-Pro.local",
				  "name" : "Yurys-MacBook-Pro.local"
				},
				"type" : "log",
				"fields" : null,
				"host" : "Yurys-MacBook-Pro.local",
				"tags" : [ "beats_input_codec_plain_applied" ],
				"app" : "alice"
			      },
			      "sort" : [ 128 ]
			    }, {
			      "_index" : "logstash-2016.07.25",
			      "_type" : "log",
			      "_id" : "AVYkNv542Gim_t2htKPR",
			      "_score" : null,
			      "_source" : {
				"message" : "Alice message 8",
				"@version" : "1",
				"@timestamp" : "2016-07-25T22:39:55.760Z",
				"type" : "log",
				"input_type" : "log",
				"source" : "/Users/yury/logs/alice.log",
				"count" : 1,
				"fields" : null,
				"beat" : {
				  "hostname" : "Yurys-MacBook-Pro.local",
				  "name" : "Yurys-MacBook-Pro.local"
				},
				"offset" : 112,
				"host" : "Yurys-MacBook-Pro.local",
				"tags" : [ "beats_input_codec_plain_applied" ],
				"app" : "alice"
			      },
			      "sort" : [ 112 ]
			    } ]
			  }
			}`

		w.Write([]byte(resp))
	}

	s, err := MockService(ts.URL)
	assert.NoError(t, err)

	expectedMessages := []string{
		"Alice message 8",
		"Alice message 9",
		"Alice message 10",
	}
	actualMessages, err := s.GetLog("app", 3)
	assert.NoError(t, err)
	assert.Equal(t, expectedMessages, actualMessages)
}

func MockService(url string) (Service, error) {
	client, err := elastic.NewSimpleClient(elastic.SetURL(url))
	if err != nil {
		return nil, err
	}
	return &service{elasticClient: client}, nil
}

Instead of hardcoding response, you can read it from a file.

By the way, if you are not sure why resp is casted to []byte in w.Write([]byte(resp)), have a look at How To Correctly Serialize JSON String In Golang, which describes the issue in details.

A couple of words on integration testing

Even though the post is devoted to unit testing Go code calling external service, I must admit that a better way to test such code would be via integration testing. Integration tests are executed on the whole system with all its components up and running. This makes the system as close to production as possible and testing the whole system provides much higher quality guarantees.

In practice, however, integration tests are usually more difficult to implement and require much more time to run. As result, integration tests are written less often than they should be.

The bottom line

When unit testing Go code calling external RESTful service via Go client, the easiest approach is to mock the client. If Go client is represented by struct (instead of interface), you can mock RESTful API itself and make the mock to serve pre-recorded real responses.