handler_test.go

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/v2/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)
}