package telemetry
import (
"bytes"
"context"
"fmt"
"io"
"net/http"
"net/http/httptest"
"reflect"
"testing"
"time"
"github.com/gogo/protobuf/proto"
pr "github.com/influxdata/influxdb/prometheus"
dto "github.com/prometheus/client_model/go"
"github.com/prometheus/common/expfmt"
"go.uber.org/zap/zaptest"
)
func TestPushGateway_Handler(t *testing.T) {
type fields struct {
Store *mockStore
now func() time.Time
}
type args struct {
w *httptest.ResponseRecorder
r *http.Request
}
tests := []struct {
name string
fields fields
args args
contentType string
wantStatus int
want []byte
}{
{
name: "unknown content-type is a bad request",
fields: fields{
Store: &mockStore{},
},
args: args{
w: httptest.NewRecorder(),
r: httptest.NewRequest("POST", "/", nil),
},
wantStatus: http.StatusBadRequest,
},
{
name: "bad metric with timestamp is a bad request",
fields: fields{
Store: &mockStore{},
now: func() time.Time { return time.Unix(0, 0) },
},
args: args{
w: httptest.NewRecorder(),
r: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{badMetric()},
),
),
},
contentType: string(expfmt.FmtProtoDelim),
wantStatus: http.StatusBadRequest,
},
{
name: "store error is an internal server error",
fields: fields{
Store: &mockStore{
err: fmt.Errorf("e1"),
},
now: func() time.Time { return time.Unix(0, 0) },
},
args: args{
w: httptest.NewRecorder(),
r: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{NewCounter("mf1", 1.0, pr.L("n1", "v1"))},
),
),
},
contentType: string(expfmt.FmtProtoDelim),
wantStatus: http.StatusInternalServerError,
want: []byte(`[{"name":"mf1","type":0,"metric":[{"label":[{"name":"n1","value":"v1"}],"counter":{"value":1},"timestamp_ms":0}]}]`),
},
{
name: "metric store in store",
fields: fields{
Store: &mockStore{},
now: func() time.Time { return time.Unix(0, 0) },
},
args: args{
w: httptest.NewRecorder(),
r: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{NewCounter("mf1", 1.0, pr.L("n1", "v1"))},
),
),
},
contentType: string(expfmt.FmtProtoDelim),
wantStatus: http.StatusAccepted,
want: []byte(`[{"name":"mf1","type":0,"metric":[{"label":[{"name":"n1","value":"v1"}],"counter":{"value":1},"timestamp_ms":0}]}]`),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
p := NewPushGateway(
zaptest.NewLogger(t),
tt.fields.Store,
&AddTimestamps{
now: tt.fields.now,
},
)
p.Encoder = &pr.JSON{}
tt.args.r.Header.Set("Content-Type", tt.contentType)
p.Handler(tt.args.w, tt.args.r)
if tt.args.w.Code != http.StatusAccepted {
t.Logf("Body: %s", tt.args.w.Body.String())
}
if got, want := tt.args.w.Code, tt.wantStatus; got != want {
t.Errorf("PushGateway.Handler() StatusCode = %v, want %v", got, want)
}
if got, want := tt.fields.Store.data, tt.want; string(got) != string(want) {
t.Errorf("PushGateway.Handler() Data = %s, want %s", got, want)
}
})
}
}
func Test_decodePostMetricsRequest(t *testing.T) {
type args struct {
req *http.Request
maxBytes int64
}
tests := []struct {
name string
args args
contentType string
want []*dto.MetricFamily
wantErr bool
}{
{
name: "bad body returns no metrics",
args: args{
req: httptest.NewRequest("POST", "/", bytes.NewBuffer([]byte{0x10})),
maxBytes: 10,
},
contentType: string(expfmt.FmtProtoDelim),
want: []*dto.MetricFamily{},
},
{
name: "no body returns no metrics",
args: args{
req: httptest.NewRequest("POST", "/", nil),
maxBytes: 10,
},
contentType: string(expfmt.FmtProtoDelim),
want: []*dto.MetricFamily{},
},
{
name: "metrics are returned from POST",
args: args{
req: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{NewCounter("mf1", 1.0, pr.L("n1", "v1"))},
),
),
maxBytes: 31,
},
contentType: string(expfmt.FmtProtoDelim),
want: []*dto.MetricFamily{NewCounter("mf1", 1.0, pr.L("n1", "v1"))},
},
{
name: "max bytes limits on record boundary returns a single record",
args: args{
req: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{
NewCounter("mf1", 1.0, pr.L("n1", "v1")),
NewCounter("mf2", 1.0, pr.L("n2", "v2")),
},
),
),
maxBytes: 31,
},
contentType: string(expfmt.FmtProtoDelim),
want: []*dto.MetricFamily{NewCounter("mf1", 1.0, pr.L("n1", "v1"))},
},
{
name: "exceeding max bytes returns an error",
args: args{
req: httptest.NewRequest("POST", "/",
mustEncode(t,
[]*dto.MetricFamily{
NewCounter("mf1", 1.0, pr.L("n1", "v1")),
NewCounter("mf2", 1.0, pr.L("n2", "v2")),
},
),
),
maxBytes: 33,
},
contentType: string(expfmt.FmtProtoDelim),
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tt.args.req.Header.Set("Content-Type", tt.contentType)
got, err := decodePostMetricsRequest(tt.args.req.Body, expfmt.Format(tt.contentType), tt.args.maxBytes)
if (err != nil) != tt.wantErr {
t.Errorf("decodePostMetricsRequest() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !reflect.DeepEqual(got, tt.want) {
t.Errorf("decodePostMetricsRequest() = %v, want %v", got, tt.want)
}
})
}
}
func badMetric() *dto.MetricFamily {
return &dto.MetricFamily{
Name: proto.String("bad"),
Type: dto.MetricType_COUNTER.Enum(),
Metric: []*dto.Metric{
&dto.Metric{
Label: []*dto.LabelPair{pr.L("n1", "v1")},
Counter: &dto.Counter{
Value: proto.Float64(1.0),
},
TimestampMs: proto.Int64(1),
},
},
}
}
func goodMetric() *dto.MetricFamily {
return &dto.MetricFamily{
Name: proto.String("good"),
Type: dto.MetricType_COUNTER.Enum(),
Metric: []*dto.Metric{
&dto.Metric{
Label: []*dto.LabelPair{pr.L("n1", "v1")},
Counter: &dto.Counter{
Value: proto.Float64(1.0),
},
},
},
}
}
func Test_valid(t *testing.T) {
type args struct {
mfs []*dto.MetricFamily
}
tests := []struct {
name string
args args
wantErr bool
}{
{
name: "metric with timestamp is invalid",
args: args{
mfs: []*dto.MetricFamily{badMetric()},
},
wantErr: true,
},
{
name: "metric without timestamp is valid",
args: args{
mfs: []*dto.MetricFamily{goodMetric()},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if err := valid(tt.args.mfs); (err != nil) != tt.wantErr {
t.Errorf("valid() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
type mockStore struct {
data []byte
err error
}
func (m *mockStore) WriteMessage(ctx context.Context, data []byte) error {
m.data = data
return m.err
}
func mustEncode(t *testing.T, mfs []*dto.MetricFamily) io.Reader {
b, err := pr.EncodeExpfmt(mfs)
if err != nil {
t.Fatalf("unable to encode %v", err)
}
return bytes.NewBuffer(b)
}