使用httptest 实现golang ES Mock
Contents
sqlmoc
在日常开发中，为了保证代码健壮性，和避免手残写错ORM语句，我们难免会遇到mock 测试，在使用MySQL 或者pg 的时候，我们可以使用开源的drive和mock 工具进行快速mock 比如 go-sqlmock，但是如果想对ES 生成对DDL语句做校验的时候，有时候显得没有sqlmock 这么方便，于是就简单的写了一个es 的mock。
思路
我连接ES 使用的是这个第三方包 github.com/olivere/elastic/v7。ES查询实际上就是调用HTTP请求，那么问题就迎刃而解了，我们只需要做一个HTTP mock 就可以解决我们的问题了。
重写相关方法
好在golang 内置的包已经提供了测试 http的包 httptest,接着我们只需要做亿点点改造就可以成为我们的es mock。
首先定义一个struct
type MockHttpClient struct {
RequestBody []byte // except genator ddl
ResponseBody []byte // except response
URL string // connect url
}
因为要实现一个http 的clinet，查看文档，控制请求及响应的方法是client Transport 方法，这个方法的数据类型是RoundTripper这个interface，我们只需要实现 RoundTrip这个方法，就可以控制http 的请求及响应了。当然我这里只是简单的实现, 真正的mock比这复杂的多。
func (c *MockHttpClient) RoundTrip(req *http.Request) (*http.Response, error) {
recorder := httptest.NewRecorder()
if c.ResponseBody == nil || c.RequestBody == nil {
return nil, errors.New("ResponseBody or RequestBody can't be null")
}
recorder.Header().Set("Content-Type", "application/json")
reqBody, err := ioutil.ReadAll(req.Body)
if err != nil {
return nil, err
}
reqReg := regexp.MustCompile(string(c.RequestBody))
regRes := reqReg.Find(reqBody)
if string(regRes) == "" {
return nil, errors.Errorf(`Except query string: "%s" But get "%s"`, reqBody, c.RequestBody)
}
if _, err := recorder.Write(c.ResponseBody); err != nil {
return nil, err
}
return recorder.Result(), nil
}
生成ES client
http client 的问题已经解决了，接下来就是用这个假的http client 生成一个es 的client
func (c *MockHttpClient) MockElasticSearchClient() (*elastic.Client, error) {
client, err := elastic.NewClient(
elastic.SetURL(c.URL),
elastic.SetSniff(false),
elastic.SetHealthcheck(false),
elastic.SetHttpClient(&http.Client{Transport: c}), elastic.SetTraceLog(logger.NewLogger("ESMock")))
return client, err
}
再加亿点点细节
新增一个 初始化mock 的方法
func NewMock() (httpMock *MockHttpClient) {
testServer := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
defer testServer.Close()
var mock = &MockHttpClient{
URL: testServer.URL,
}
return mock
}
为 mock client 提供一个 设置request body 和response body 的方法
func (c *MockHttpClient) ShouldRequestQuery(queryString string) *MockHttpClient {
c.RequestBody = []byte(queryString)
return c
}
func (c *MockHttpClient) ShouldResponse(Response string) *MockHttpClient {
c.ResponseBody = []byte(Response)
return c
}
最后的调用
func TestRepository_AlarmCount(t *testing.T) {
mock := esMock.NewMock()
// 设置预期的值
mock.ShouldRequestQuery(exceptQuery) // exceptQuery 为预期产生的DDL Query
mock.ShouldResponse(exceptResp) //exceptResp 为预期的响应
client, err := MockElasticSearchClient()
// 我们项目把数据请求都封再repository，也可以直接掉client 的search 等方法
repo, err := esRepository.NewAttackPhaseRepository(client)
assert.NoError(t, err)
actualRes, err := repo.AlarmCount(ctx, &esRepository.AlarmFilter{})
assert.NoError(t, err)
assert.Equal(t, len(actualRes), 4)
// 没封装查询方法可以直接写成这样 client.Search("index_name")...
}