consul/agent/rpc/middleware/interceptors_test.go

package middleware

import (
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/armon/go-metrics"
	"github.com/hashicorp/go-hclog"
	"github.com/stretchr/testify/require"
)

// obs holds all the things we want to assert on that we recorded correctly in our tests.
type obs struct {
	key     []string
	elapsed float32
	labels  []metrics.Label
}

// recorderStore acts as an in-mem mock storage for all the RequestRecorder.Record() RecorderFunc calls.
type recorderStore struct {
	lock  sync.Mutex
	store map[string]obs
}

func (rs *recorderStore) put(key []string, o obs) {
	rs.lock.Lock()
	defer rs.lock.Unlock()

	actualKey := strings.Join(append(key, o.labels[0].Value), "")
	rs.store[actualKey] = o
}

func (rs *recorderStore) get(key []string) obs {
	rs.lock.Lock()
	defer rs.lock.Unlock()

	actualKey := strings.Join(key, "")
	return rs.store[actualKey]
}

var store = recorderStore{store: make(map[string]obs)}
var simpleRecorderFunc = func(key []string, val float32, labels []metrics.Label) {
	o := obs{key: key, elapsed: val, labels: labels}
	store.put(key, o)
}

type readRequest struct{}
type writeRequest struct{}

func (rr readRequest) IsRead() bool {
	return true
}

func (wr writeRequest) IsRead() bool {
	return false
}

// TestRequestRecorder_SimpleOK tests that the RequestRecorder can record a simple request.
func TestRequestRecorder_SimpleOK(t *testing.T) {
	t.Parallel()

	r := RequestRecorder{
		Logger:       hclog.NewInterceptLogger(&hclog.LoggerOptions{}),
		RecorderFunc: simpleRecorderFunc,
	}

	start := time.Now()
	r.Record("A.B", RPCTypeInternal, start, struct{}{}, false)

	expectedLabels := []metrics.Label{
		{Name: "method", Value: "A.B"},
		{Name: "errored", Value: "false"},
		{Name: "request_type", Value: "unreported"},
		{Name: "rpc_type", Value: RPCTypeInternal},
	}

	o := store.get(append(metricRPCRequest, expectedLabels[0].Value))
	require.Equal(t, o.key, metricRPCRequest)
	require.LessOrEqual(t, o.elapsed, float32(start.Sub(time.Now()).Milliseconds()))
	require.Equal(t, o.labels, expectedLabels)
}

// TestRequestRecorder_ReadRequest tests that RequestRecorder can record a read request AND a responseErrored arg.
func TestRequestRecorder_ReadRequest(t *testing.T) {
	t.Parallel()

	r := RequestRecorder{
		Logger:       hclog.NewInterceptLogger(&hclog.LoggerOptions{}),
		RecorderFunc: simpleRecorderFunc,
	}

	start := time.Now()

	r.Record("B.A", RPCTypeNetRPC, start, readRequest{}, true)

	expectedLabels := []metrics.Label{
		{Name: "method", Value: "B.A"},
		{Name: "errored", Value: "true"},
		{Name: "request_type", Value: "read"},
		{Name: "rpc_type", Value: RPCTypeNetRPC},
	}

	o := store.get(append(metricRPCRequest, expectedLabels[0].Value))
	require.Equal(t, o.labels, expectedLabels)
}

// TestRequestRecorder_WriteRequest tests that RequestRecorder can record a write request.
func TestRequestRecorder_WriteRequest(t *testing.T) {
	t.Parallel()

	r := RequestRecorder{
		Logger:       hclog.NewInterceptLogger(&hclog.LoggerOptions{}),
		RecorderFunc: simpleRecorderFunc,
	}

	start := time.Now()

	r.Record("B.C", RPCTypeNetRPC, start, writeRequest{}, true)

	expectedLabels := []metrics.Label{
		{Name: "method", Value: "B.C"},
		{Name: "errored", Value: "true"},
		{Name: "request_type", Value: "write"},
		{Name: "rpc_type", Value: RPCTypeNetRPC},
	}

	o := store.get(append(metricRPCRequest, expectedLabels[0].Value))
	require.Equal(t, o.labels, expectedLabels)
}