connect-go の interceptor に外部から context を渡す

概要

connect-go の interceptor に外部から context を渡す方法についてです。

あまりすることは無さそうですが、context に値をセットするような interceptor を書いていて、単体テストを書く際に今回のような方法が有用だったので紹介します。

interceptor

package sample

import (
    "context"

    "github.com/bufbuild/connect-go"
)

const CtxKey = "key"

func NewSampleInterceptor() connect.UnaryInterceptorFunc {
    interceptor := func(next connect.UnaryFunc) connect.UnaryFunc {
        return connect.UnaryFunc(func(
            ctx context.Context,
            req connect.AnyRequest,
        ) (connect.AnyResponse, error) {
            ctx = context.WithValue(ctx, CtxKey, "hoge")
            return next(ctx, req)
        })
    }
    return connect.UnaryInterceptorFunc(interceptor)
}

サンプルは context に値をセットするだけの interceptor です。
単体テストを書く場合 context.WithValue() が期待通りに呼ばれたことでも良い気がしますが、今回は context の値を比較して検証したいと思いました。

この関数は高階関数になっており、そのことを正しく理解していれば上手くテストを書くことができたのですが理解不足のため context の受け渡しに悩みました。

test

package sample_test

import (
    "context"
    "sample/cmd/server/interceptor/sample"
    "testing"

    "github.com/bufbuild/connect-go"
    "github.com/stretchr/testify/assert"
)

func TestNewSampleInterceptor(t *testing.T) {
    interceptor := sample.NewSampleInterceptor()

    ctx := context.Background()
    req := connect.NewRequest(&struct{}{})

    next := connect.UnaryFunc(func(
        ctx context.Context,
        req connect.AnyRequest,
    ) (connect.AnyResponse, error) {
        assert.Equal(t, "hoge", ctx.Value(sample.CtxKey))
        return nil, nil
    })
    interceptor(next)(ctx, req)
}

結論このようにテストを書きました。

変数 interceptor は関数型の UnaryInterceptorFunc が入り、引数の型は UnaryFunc です。

そのため新たに定義した UnaryFunc (変数 next) 内で値の確認を行うようにしています。

これで NewSampleInterceptor 内で定義している関数に ctx, req を渡し、引数で渡している関数 next に上手く ctx, req を渡して context の検証をすることができます。

まとめ

Go における高階関数について理解不足だったので悩みましたが、なんとかやりたいことはできました。
もっと良いやり方があればぜひ教えてください！