Add client

Signed-off-by: Jorge Turrado <[email protected]>

.github/workflows/build.yaml

@@ -0,0 +1,25 @@
+name: Build
+
+on:
+  push:
+    branches: 
+    - main
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Go
+      uses: actions/setup-go@v3
+      with:
+        go-version: "1.20"
+
+    - name: Install XK6
+      run:  go install go.k6.io/xk6/cmd/xk6@latest
+
+    - name: Test
+      run: make test

Makefile

@@ -0,0 +1,5 @@
+build:
+	xk6 build --with github.com/JorTurFer/xk6-input-prometheus=.
+
+test: build
+	./k6 run -i 1 -u 1 example.js

example.js

@@ -0,0 +1,7 @@
+import prometheus from 'k6/x/prometheusread';
+
+export default function () {
+  var client = prometheus.newPrometheusClient("http://demo.robustperception.io:9090", "user", "password")
+  var response = client.queryScalar("sum(alertmanager_notifications_total)")
+  console.log(response)
+}

go.mod

@@ -0,0 +1,51 @@
+module github.com/JorTurFer/xk6-input-prometheus
+
+go 1.20
+
+require (
+	github.com/prometheus/client_golang v1.16.0
+	github.com/prometheus/common v0.42.0
+	github.com/stretchr/testify v1.8.2
+	go.k6.io/k6 v0.45.0
+)
+
+require (
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dlclark/regexp2 v1.9.0 // indirect
+	github.com/dop251/goja v0.0.0-20230531210528-d7324b2d74f7 // indirect
+	github.com/fatih/color v1.15.0 // indirect
+	github.com/go-sourcemap/sourcemap v2.1.4-0.20211119122758-180fcef48034+incompatible // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/google/pprof v0.0.0-20230207041349-798e818bf904 // indirect
+	github.com/josharian/intern v1.0.0 // indirect
+	github.com/jpillora/backoff v1.0.0 // indirect
+	github.com/json-iterator/go v1.1.12 // indirect
+	github.com/mailru/easyjson v0.7.7 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.18 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
+	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
+	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/mstoykov/atlas v0.0.0-20220811071828-388f114305dd // indirect
+	github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
+	github.com/onsi/ginkgo v1.16.5 // indirect
+	github.com/onsi/gomega v1.27.10 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_model v0.3.0 // indirect
+	github.com/prometheus/procfs v0.10.1 // indirect
+	github.com/serenize/snaker v0.0.0-20201027110005-a7ad2135616e // indirect
+	github.com/sirupsen/logrus v1.9.0 // indirect
+	github.com/spf13/afero v1.1.2 // indirect
+	golang.org/x/net v0.12.0 // indirect
+	golang.org/x/oauth2 v0.5.0 // indirect
+	golang.org/x/sys v0.10.0 // indirect
+	golang.org/x/text v0.11.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/protobuf v1.30.0 // indirect
+	gopkg.in/guregu/null.v3 v3.3.0 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

input_prometheus.go

@@ -0,0 +1,75 @@
+package prometheus
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/prometheus/client_golang/api"
+	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
+	"github.com/prometheus/common/config"
+	"github.com/prometheus/common/model"
+	"go.k6.io/k6/js/modules"
+)
+
+// init is called by the Go runtime at application startup.
+func init() {
+	modules.Register("k6/x/prometheusread", new(Prometheus))
+}
+
+type Prometheus struct{}
+
+type Client struct {
+	url      string
+	username string
+	password config.Secret
+}
+
+func (*Prometheus) NewPrometheusClient(url, username string, password config.Secret) Client {
+	return Client{
+		url:      url,
+		username: username,
+		password: password,
+	}
+}
+
+func (c *Client) QueryScalar(query string) (int64, error) {
+	client, err := c.generateClient()
+	if err != nil {
+		return -1, err
+	}
+	v1api := v1.NewAPI(client)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	result, warnings, err := v1api.Query(ctx, query, time.Now(), v1.WithTimeout(5*time.Second))
+	if err != nil {
+		return -1, err
+	}
+	if len(warnings) > 0 {
+		return -1, err
+	}
+
+	value := 0
+	switch obj := result.(type) {
+	case model.Vector:
+		if obj.Len() == 0 {
+			return -1, errors.New("empty vector")
+		}
+		sample := obj[len(obj)-1]
+		value = int(sample.Value)
+	case *model.Scalar:
+		value = int(obj.Value)
+	}
+	return int64(value), nil
+}
+
+func (c *Client) generateClient() (api.Client, error) {
+	roundTripper := api.DefaultRoundTripper
+	if c.password != "" {
+		roundTripper = config.NewBasicAuthRoundTripper(c.username, c.password, "", api.DefaultRoundTripper)
+	}
+	return api.NewClient(api.Config{
+		Address:      c.url,
+		RoundTripper: roundTripper,
+	})
+}

input_prometheus_test.go

@@ -0,0 +1,43 @@
+package prometheus_test
+
+import (
+	"testing"
+
+	prometheus "github.com/JorTurFer/xk6-input-prometheus"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestQueryScalar(t *testing.T) {
+	module := prometheus.Prometheus{}
+	client := module.NewPrometheusClient("http://demo.robustperception.io:9090", "", "")
+
+	result, err := client.QueryScalar("up")
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, result, int64(1))
+}
+
+func TestQueryScalarWithOperation(t *testing.T) {
+	module := prometheus.Prometheus{}
+	client := module.NewPrometheusClient("http://demo.robustperception.io:9090", "", "")
+
+	result, err := client.QueryScalar("sum(alertmanager_notifications_total)")
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, result, int64(1))
+}
+
+func TestQueryScalarWithBasicAuth(t *testing.T) {
+	module := prometheus.Prometheus{}
+	client := module.NewPrometheusClient("http://demo.robustperception.io:9090", "test", "1234")
+
+	result, err := client.QueryScalar("sum(alertmanager_notifications_total)")
+	assert.NoError(t, err)
+	assert.GreaterOrEqual(t, result, int64(1))
+}
+
+func TestQueryScalarWithError(t *testing.T) {
+	module := prometheus.Prometheus{}
+	client := module.NewPrometheusClient("http://demo.robustperception.io:9090", "test", "1234")
+
+	_, err := client.QueryScalar("sum(random_metric)")
+	assert.Error(t, err)
+}