147 lines
3.6 KiB
Go
147 lines
3.6 KiB
Go
package query_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"io"
|
|
"io/ioutil"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/go-cmp/cmp/cmpopts"
|
|
"github.com/influxdata/flux"
|
|
platform "github.com/influxdata/influxdb/v2"
|
|
"github.com/influxdata/influxdb/v2/query"
|
|
"github.com/influxdata/influxdb/v2/query/mock"
|
|
"github.com/opentracing/opentracing-go"
|
|
"github.com/uber/jaeger-client-go"
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
var orgID = MustIDBase16("ba55ba55ba55ba55")
|
|
|
|
// MustIDBase16 is an helper to ensure a correct ID is built during testing.
|
|
func MustIDBase16(s string) platform.ID {
|
|
id, err := platform.IDFromString(s)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return *id
|
|
}
|
|
|
|
var opts = []cmp.Option{
|
|
cmpopts.IgnoreUnexported(query.ProxyRequest{}),
|
|
cmpopts.IgnoreUnexported(query.Request{}),
|
|
}
|
|
|
|
func TestLoggingProxyQueryService(t *testing.T) {
|
|
// Set a Jaeger in-memory tracer to get span information in the query log.
|
|
oldTracer := opentracing.GlobalTracer()
|
|
defer opentracing.SetGlobalTracer(oldTracer)
|
|
sampler := jaeger.NewConstSampler(true)
|
|
reporter := jaeger.NewInMemoryReporter()
|
|
tracer, closer := jaeger.NewTracer(t.Name(), sampler, reporter)
|
|
defer closer.Close()
|
|
opentracing.SetGlobalTracer(tracer)
|
|
|
|
wantStats := flux.Statistics{
|
|
TotalDuration: time.Second,
|
|
CompileDuration: time.Second,
|
|
QueueDuration: time.Second,
|
|
PlanDuration: time.Second,
|
|
RequeueDuration: time.Second,
|
|
ExecuteDuration: time.Second,
|
|
Concurrency: 2,
|
|
MaxAllocated: 2048,
|
|
}
|
|
wantBytes := 10
|
|
pqs := &mock.ProxyQueryService{
|
|
QueryF: func(ctx context.Context, w io.Writer, req *query.ProxyRequest) (flux.Statistics, error) {
|
|
w.Write(make([]byte, wantBytes))
|
|
return wantStats, nil
|
|
},
|
|
}
|
|
var logs []query.Log
|
|
logger := &mock.QueryLogger{
|
|
LogFn: func(l query.Log) error {
|
|
logs = append(logs, l)
|
|
return nil
|
|
},
|
|
}
|
|
|
|
req := &query.ProxyRequest{
|
|
Request: query.Request{
|
|
Authorization: nil,
|
|
OrganizationID: orgID,
|
|
Compiler: nil,
|
|
},
|
|
Dialect: nil,
|
|
}
|
|
|
|
t.Run("log", func(t *testing.T) {
|
|
defer func() {
|
|
logs = nil
|
|
}()
|
|
wantTime := time.Now()
|
|
lpqs := query.NewLoggingProxyQueryService(zap.NewNop(), logger, pqs)
|
|
lpqs.SetNowFunctionForTesting(func() time.Time {
|
|
return wantTime
|
|
})
|
|
|
|
var buf bytes.Buffer
|
|
stats, err := lpqs.Query(context.Background(), &buf, req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !cmp.Equal(wantStats, stats, opts...) {
|
|
t.Errorf("unexpected query stats: -want/+got\n%s", cmp.Diff(wantStats, stats, opts...))
|
|
}
|
|
traceID := reporter.GetSpans()[0].Context().(jaeger.SpanContext).TraceID().String()
|
|
wantLogs := []query.Log{{
|
|
Time: wantTime,
|
|
OrganizationID: orgID,
|
|
TraceID: traceID,
|
|
Sampled: true,
|
|
Error: nil,
|
|
ProxyRequest: req,
|
|
ResponseSize: int64(wantBytes),
|
|
Statistics: wantStats,
|
|
}}
|
|
if !cmp.Equal(wantLogs, logs, opts...) {
|
|
t.Errorf("unexpected query logs: -want/+got\n%s", cmp.Diff(wantLogs, logs, opts...))
|
|
}
|
|
})
|
|
|
|
t.Run("conditional logging", func(t *testing.T) {
|
|
defer func() {
|
|
logs = nil
|
|
}()
|
|
|
|
loggingKey := "do-logging"
|
|
condLog := query.ConditionalLogging(func(ctx context.Context) bool {
|
|
return ctx.Value(loggingKey) != nil
|
|
})
|
|
|
|
lpqs := query.NewLoggingProxyQueryService(zap.NewNop(), logger, pqs, condLog)
|
|
_, err := lpqs.Query(context.Background(), ioutil.Discard, req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if len(logs) != 0 {
|
|
t.Fatal("expected query service not to log")
|
|
}
|
|
|
|
ctx := context.WithValue(context.Background(), loggingKey, true)
|
|
_, err = lpqs.Query(ctx, ioutil.Discard, req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
if len(logs) != 1 {
|
|
t.Fatal("expected query service to log")
|
|
}
|
|
})
|
|
}
|