127 lines
3.3 KiB
Go
127 lines
3.3 KiB
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
|
|
start time.Time
|
|
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, start time.Time, labels []metrics.Label) {
|
|
o := obs{key: key, start: start, 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) {
|
|
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: "write"},
|
|
{Name: "rpc_type", Value: RPCTypeInternal},
|
|
}
|
|
|
|
o := store.get(append(metricRPCRequest, expectedLabels[0].Value))
|
|
require.Equal(t, o.key, metricRPCRequest)
|
|
require.Equal(t, o.start, start)
|
|
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) {
|
|
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) {
|
|
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)
|
|
}
|