emahiro/b.log

Drastically Repeat Yourself !!!!

トップ > go > cloud.google.com/go/bigquery.Query の結果を go test の中で mock する

cloud.google.com/go/bigquery.Query の結果を go test の中で mock する

go bigquery

Overview

BigQuery を使った実装をしているときに、API リクエストを mock したいケースがあると思いますが、そういうときに安易に DI を使わずに実装する方法を考えます

やること

やることは難しくなくて BigQuery のクライアントを生成するときにテストでは API のリクエスト先を httptest.Server に向けるだけです。

実際のクライアントを生成する実装は以下のようになります。

ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    fmt.Printf("Request URL: %v\n", r.URL.String())
    var res []byte
    if _, err := w.Write(res); err != nil {...}
})

bq, err := bigquery.NewClient(ctx, "$ProjectID", option.WithEndpoint(ts.URL)
if err != nil {
    return err
}

httptest.Server を起動し、BigQuery の接続先をOptionで起動したテストサーバーに向けます。

Query Operation の場合

https://pkg.go.dev/cloud.google.com/go/bigquery@v1.59.1#Client.Query のオペレーションを例にサンプルでテストサーバーを実装します。

Query Operation は BigQuery ライブラリの中で job のリストを取得しており、当てはまる jobId を query の API をリクエストする、という振る舞いになっています。

ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    fmt.Printf("Request URL: %v\n", r.URL.String())
    var res []byte

    if r.URL.Path == "/projects/eure-metis-dev/jobs" {
    res = []byte(`{"kind":"bigquery#job","etag":"\"abcdefghijklmnopqrstuvwxyz\"","id":"your-project-id:your-job-id","selfLink":"https://www.googleapis.com/bigquery/v2/projects/your-project-id/jobs/your-job-id","jobReference":{"projectId":"your-project-id","jobId":"your-job-id","location":"US"},"configuration":{"query":{"query":"SELECT user_id FROM test-data","destinationTable":{"projectId":"your-project-id","datasetId":"your-dataset-id","tableId":"your-table-id"},"writeDisposition":"WRITE_TRUNCATE","useQueryCache":true}},"status":{"state":"DONE","errorResult":null,"errors":null},"statistics":{"creationTime":"1600000000000","startTime":"1600000001000","endTime":"1600000005000","totalBytesProcessed":"123456","query":{"totalBytesBilled":"123456","cacheHit":false,"totalBytesProcessed":"123456","totalSlotMs":"1234","numDmlAffectedRows":"0"}},"user_email":"user@example.com"}`)}

    if r.URL.Path == "/projects/your-project-id/queries/your-job-id" {
    res = []byte(`{"kind":"bigquery#getQueryResultsResponse","etag":"\"abcdefghijklmn\"","schema":{"fields":[{"name":"user_id","type":"INTEGER","mode":"NULLABLE"}]},"jobReference":{"projectId":"your-project-id","jobId":"your-job-id"},"totalRows":"10","rows":[{"f":[{"v":"1"}]}],"totalBytesProcessed":"123456","jobComplete":true,"cacheHit":false}`)}

    if _, err := w.Write(res); err != nil {...}
})

bq, err := bigquery.NewClient(ctx, "$ProjectID", option.WithEndpoint(ts.URL)
if err != nil {
    return err
}

これで２つ目の query のレスポンスに入っている、field key の結果が返ってきます。

ハマったところ

BigQuery のライブラリでどの API を Call しているのかわからなかった。

API のリクエストとレスポンスを調べて httptest.Server の中でレスポンスを生成して w.Write すれば終わりかと思っていたのですが、BigQuery のライブラリでは job の一覧を call したあと、該当する job の query を取得する API をリクエストしてることがわかりました。

これは httptest.Server で起動したサーバーでリクエストを dump してみて、job list -> query と順番に Call してることがわかりました。これで httptest.Server の中で call された URL に応じてレスポンスを変えることで、BigQuery の API ライブラリの振る舞いをトレースすることができるようになりました。

BigQuery の API のレスポンスを生成する

API のドキュメントは以下の２つを読めば Query オペレーションの振る舞いはトレースできるようになります。

ただし、この API ドキュメントもそんなに優しいわけではなく、例えば struct みたいな感じで中身が具体的に分からないフィールドもあるのに、API のレスポンスが期待されるレスポンスでないと空のレスポンスを作っているにも関わらず、ライブラリ側の Unmarshal 処理でエラーになっていまう、ということがありました。（レスポンスそれ自体に意味はないけど、ライブラリ側のバリデーションで落とされては振る舞いが完全なものになりません。

ちょっとどうしようか悩んでましたが、ChatGPT に上記のドキュメントを渡して「ダミーのレスポンスを生成して」という雑なプロンプトを投げるだけでサンプルレスポンスを生成してくれたのでこれはこれですぐに解決しました。

Google の API のリクエストを mock するときのサンプルレスポンスの生成は今後全部 ChatGPT に投げれば解決しそうでした。