Merge Prometheus remote write output

go.mod

@@ -16,7 +16,6 @@ require (
 	github.com/gorilla/websocket v1.5.3
 	github.com/grafana/xk6-dashboard v0.7.5
 	github.com/grafana/xk6-output-opentelemetry v0.3.0
-	github.com/grafana/xk6-output-prometheus-remote v0.5.1
 	github.com/grafana/xk6-redis v0.3.3
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/influxdata/influxdb1-client v0.0.0-20190402204710-8ff2fc3824fc
@@ -57,7 +56,7 @@ require (
 
 require (
 	buf.build/gen/go/gogo/protobuf/protocolbuffers/go v1.31.0-20210810001428-4df00b267f94.1 // indirect
-	buf.build/gen/go/prometheus/prometheus/protocolbuffers/go v1.31.0-20230627135113-9a12bc2590d2.1 // indirect
+	buf.build/gen/go/prometheus/prometheus/protocolbuffers/go v1.31.0-20230627135113-9a12bc2590d2.1
 	github.com/andybalholm/cascadia v1.3.2 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bufbuild/protocompile v0.14.1 // indirect
@@ -80,8 +79,8 @@ require (
 	github.com/nxadm/tail v1.4.11 // indirect
 	github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/client_golang v1.16.0 // indirect
-	github.com/prometheus/client_model v0.4.0 // indirect
+	github.com/prometheus/client_golang v1.16.0
+	github.com/prometheus/client_model v0.4.0
 	github.com/prometheus/common v0.42.0 // indirect
 	github.com/prometheus/procfs v0.10.1 // indirect
 	github.com/r3labs/sse/v2 v2.10.0 // indirect

go.sum

@@ -73,8 +73,6 @@ github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaW
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
-github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
-github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
@@ -91,8 +89,6 @@ github.com/grafana/xk6-dashboard v0.7.5 h1:TcILyffT/Ea/XD7xG1jMA5lwtusOPRbEQsQDH
 github.com/grafana/xk6-dashboard v0.7.5/go.mod h1:Y75F8xmgCraKT+pBzFH6me9AyH5PkXD+Bxo1dm6Il/M=
 github.com/grafana/xk6-output-opentelemetry v0.3.0 h1:dmclGBFtFVRJijqLncpu2dKVIIvx1GS3y6sQGg4Khl8=
 github.com/grafana/xk6-output-opentelemetry v0.3.0/go.mod h1:5uUg0J2dxrw+DkC9fCsfiKv2jTpccn7nz6vGmRUr4f8=
-github.com/grafana/xk6-output-prometheus-remote v0.5.1 h1:zsB0x/1kuArqnDNh419UWdMUYZwrzqtfjTUswWPIHhI=
-github.com/grafana/xk6-output-prometheus-remote v0.5.1/go.mod h1:esXXthLoVp9JUdGkECRthESVYu0TQTR24wrx2nRM9ak=
 github.com/grafana/xk6-redis v0.3.3 h1:9QiS4OUjYMzvriGzbqrhCn9i/kENmCkibZE+xCJSdPk=
 github.com/grafana/xk6-redis v0.3.3/go.mod h1:YL1qfBZXW4o2aAF6/gRa/N9oGoumUg7cJzzCHLxoCOw=
 github.com/grpc-ecosystem/go-grpc-middleware v1.4.0 h1:UH//fgunKIs4JdUbpDl1VZCDaL56wXCB/5+wF6uHfaI=

internal/cmd/outputs.go

@@ -11,13 +11,13 @@ import (
 	"go.k6.io/k6/internal/output/cloud"
 	"go.k6.io/k6/internal/output/influxdb"
 	"go.k6.io/k6/internal/output/json"
+	"go.k6.io/k6/internal/output/prometheusrw/remotewrite"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/output"
 	"go.k6.io/k6/output/csv"
 
 	"github.com/grafana/xk6-dashboard/dashboard"
 	"github.com/grafana/xk6-output-opentelemetry/pkg/opentelemetry"
-	"github.com/grafana/xk6-output-prometheus-remote/pkg/remotewrite"
 )
 
 // builtinOutput marks the available builtin outputs.

vendor/github.com/grafana/xk6-output-prometheus-remote/pkg/remote/client.go → internal/output/prometheusrw/remote/client.go

@@ -12,7 +12,7 @@ import (
 	"net/url"
 	"time"
 
-	"github.com/grafana/xk6-output-prometheus-remote/pkg/sigv4"
+	"go.k6.io/k6/internal/output/prometheusrw/sigv4"
 
 	prompb "buf.build/gen/go/prometheus/prometheus/protocolbuffers/go"
 	"github.com/klauspost/compress/snappy"

internal/output/prometheusrw/remote/client_test.go

@@ -0,0 +1,228 @@
+package remote
+
+import (
+	"context"
+	"io"
+	"math"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	prompb "buf.build/gen/go/prometheus/prometheus/protocolbuffers/go"
+	"github.com/klauspost/compress/snappy"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.k6.io/k6/internal/output/prometheusrw/stale"
+	"google.golang.org/protobuf/proto"
+)
+
+func TestNewWriteClient(t *testing.T) {
+	t.Parallel()
+	t.Run("DefaultConfig", func(t *testing.T) {
+		t.Parallel()
+		wc, err := NewWriteClient("http://example.com/api/v1/write", nil)
+		require.NoError(t, err)
+		require.NotNil(t, wc)
+		assert.Equal(t, wc.cfg, &HTTPConfig{})
+	})
+
+	t.Run("CustomConfig", func(t *testing.T) {
+		t.Parallel()
+		hc := &HTTPConfig{Timeout: time.Second}
+		wc, err := NewWriteClient("http://example.com/api/v1/write", hc)
+		require.NoError(t, err)
+		require.NotNil(t, wc)
+		assert.Equal(t, wc.cfg, hc)
+	})
+
+	t.Run("InvalidURL", func(t *testing.T) {
+		t.Parallel()
+		wc, err := NewWriteClient("fake://bad url", nil)
+		require.Error(t, err)
+		assert.Nil(t, wc)
+	})
+}
+
+func TestClientStore(t *testing.T) {
+	t.Parallel()
+	h := func(rw http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, r.Header.Get("Content-Encoding"), "snappy")
+		assert.Equal(t, r.Header.Get("Content-Type"), "application/x-protobuf")
+		assert.Equal(t, r.Header.Get("User-Agent"), "k6-prometheus-rw-output")
+		assert.Equal(t, r.Header.Get("X-Prometheus-Remote-Write-Version"), "0.1.0")
+		assert.NotEmpty(t, r.Header.Get("Content-Length"))
+
+		b, err := io.ReadAll(r.Body)
+		assert.NoError(t, err)
+		assert.NotEmpty(t, len(b))
+
+		rw.WriteHeader(http.StatusNoContent)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{},
+	}
+	data := &prompb.TimeSeries{
+		Labels: []*prompb.Label{
+			{
+				Name:  "label1",
+				Value: "label1-val",
+			},
+		},
+		Samples: []*prompb.Sample{
+			{
+				Value:     8.5,
+				Timestamp: time.Now().UnixMilli(),
+			},
+		},
+	}
+	err = c.Store(context.Background(), []*prompb.TimeSeries{data})
+	assert.NoError(t, err)
+}
+
+func TestClientStoreHTTPError(t *testing.T) {
+	t.Parallel()
+	h := func(w http.ResponseWriter, _ *http.Request) {
+		http.Error(w, "bad", http.StatusUnauthorized)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{},
+	}
+	assert.Error(t, c.Store(context.Background(), nil))
+}
+
+func TestClientStoreHTTPBasic(t *testing.T) {
+	t.Parallel()
+	h := func(_ http.ResponseWriter, r *http.Request) {
+		u, pwd, ok := r.BasicAuth()
+		require.True(t, ok)
+		assert.Equal(t, "usertest", u)
+		assert.Equal(t, "pwdtest", pwd)
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{
+			BasicAuth: &BasicAuth{
+				Username: "usertest",
+				Password: "pwdtest",
+			},
+		},
+	}
+	assert.NoError(t, c.Store(context.Background(), nil))
+}
+
+func TestClientStoreHeaders(t *testing.T) {
+	t.Parallel()
+	h := func(_ http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, r.Header.Get("X-Prometheus-Remote-Write-Version"), "0.1.0")
+		assert.Equal(t, r.Header.Get("X-MY-CUSTOM-HEADER"), "fake")
+	}
+	ts := httptest.NewServer(http.HandlerFunc(h))
+	defer ts.Close()
+
+	u, err := url.Parse(ts.URL)
+	require.NoError(t, err)
+
+	c := &WriteClient{
+		hc:  ts.Client(),
+		url: u,
+		cfg: &HTTPConfig{
+			Headers: http.Header(map[string][]string{
+				"X-MY-CUSTOM-HEADER": {"fake"},
+				// If the same key, of a mandatory protocol's header
+				// is provided, it will be overwritten.
+				"X-Prometheus-Remote-Write-Version": {"fake"},
+			}),
+		},
+	}
+	assert.NoError(t, c.Store(context.Background(), nil))
+}
+
+func TestNewWriteRequestBody(t *testing.T) {
+	t.Parallel()
+	ts := []*prompb.TimeSeries{
+		{
+			Labels:  []*prompb.Label{{Name: "label1", Value: "val1"}},
+			Samples: []*prompb.Sample{{Value: 10.1, Timestamp: time.Unix(1, 0).Unix()}},
+		},
+	}
+	b, err := newWriteRequestBody(ts)
+	require.NoError(t, err)
+	require.NotEmpty(t, string(b))
+	assert.Contains(t, string(b), `label1`)
+}
+
+func TestNewWriteRequestBodyWithStaleMarker(t *testing.T) {
+	t.Parallel()
+
+	timestamp := time.Date(2022, time.December, 15, 11, 41, 18, 123, time.UTC)
+
+	ts := []*prompb.TimeSeries{
+		{
+			Labels: []*prompb.Label{{Name: "label1", Value: "val1"}},
+			Samples: []*prompb.Sample{{
+				Value:     stale.Marker,
+				Timestamp: timestamp.UnixMilli(),
+			}},
+		},
+	}
+	b, err := newWriteRequestBody(ts)
+	require.NoError(t, err)
+	require.NotEmpty(t, b)
+
+	sb, err := snappy.Decode(nil, b)
+	require.NoError(t, err)
+
+	var series prompb.WriteRequest
+	err = proto.Unmarshal(sb, &series)
+	require.NoError(t, err)
+	require.NotEmpty(t, series.Timeseries[0])
+	require.NotEmpty(t, series.Timeseries[0].Samples)
+
+	assert.True(t, math.IsNaN(series.Timeseries[0].Samples[0].Value))
+	assert.Equal(t, timestamp.UnixMilli(), series.Timeseries[0].Samples[0].Timestamp)
+}
+
+func TestValidateStatusCode(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		status int
+		expErr bool
+	}{
+		{status: http.StatusOK, expErr: false},        // Mimir
+		{status: http.StatusNoContent, expErr: false}, // Prometheus
+		{status: http.StatusBadRequest, expErr: true},
+	}
+	for _, tt := range tests {
+		err := validateResponseStatus(tt.status)
+		if tt.expErr {
+			assert.Error(t, err)
+			continue
+		}
+		assert.NoError(t, err)
+	}
+}

vendor/github.com/grafana/xk6-output-prometheus-remote/pkg/remotewrite/config.go → internal/output/prometheusrw/remotewrite/config.go

@@ -10,9 +10,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/grafana/xk6-output-prometheus-remote/pkg/sigv4"
+	"go.k6.io/k6/internal/output/prometheusrw/sigv4"
 
-	"github.com/grafana/xk6-output-prometheus-remote/pkg/remote"
+	"go.k6.io/k6/internal/output/prometheusrw/remote"
 	"go.k6.io/k6/lib/types"
 	"gopkg.in/guregu/null.v3"
 )