Restructuring Prometheus client, added unit-tests to regex-queries active alerts (#386)

* prometheus labels incl tests * enable label in main, add log, docs * revert the option to query by label * revert the option to query by label * PromClient instantiate by func,white space removal * revert whitespace fix for readability. * revert removal of newlines for readability * rename New to NewPromClient to improve readability Co-authored-by: simp <simp@saxobank.com>
2026-08-28 05:47:23 +00:00 · 2021-07-27 07:09:46 +02:00
parent 6e16e993d9
commit c7d5810503
5 changed files with 176 additions and 13 deletions
@@ -12,6 +12,7 @@ import (
 	"strings"
 	"time"

+	papi "github.com/prometheus/client_golang/api"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	v1 "k8s.io/api/core/v1"
@@ -228,8 +229,9 @@ type RebootBlocker interface {
 // PrometheusBlockingChecker contains info for connecting
 // to prometheus, and can give info about whether a reboot should be blocked
 type PrometheusBlockingChecker struct {
-	// URL to contact prometheus API for checking alerts
-	promURL string
+	// prometheusClient to make prometheus-go-client and api config available
+	// into the PrometheusBlockingChecker struct
+	promClient *alerts.PromClient
 	// regexp used to get alerts
 	filter *regexp.Regexp
 }
@@ -245,7 +247,8 @@ type KubernetesBlockingChecker struct {
 }

 func (pb PrometheusBlockingChecker) isBlocked() bool {
-	alertNames, err := alerts.PrometheusActiveAlerts(pb.promURL, pb.filter)
+
+	alertNames, err := pb.promClient.ActiveAlerts(pb.filter)
 	if err != nil {
 		log.Warnf("Reboot blocked: prometheus query error: %v", err)
 		return true
@@ -513,6 +516,12 @@ func rebootAsRequired(nodeID string, rebootCommand []string, sentinelCommand []s
 		preferNoScheduleTaint.Disable()
 	}

+	// instantiate prometheus client
+	promClient, err := alerts.NewPromClient(papi.Config{Address: prometheusURL})
+	if err != nil {
+		log.Fatal("Unable to create prometheus client: ", err)
+	}
+
 	source := rand.NewSource(time.Now().UnixNano())
 	tick := delaytick.New(source, period)
 	for range tick {
@@ -531,7 +540,7 @@ func rebootAsRequired(nodeID string, rebootCommand []string, sentinelCommand []s

 		var blockCheckers []RebootBlocker
 		if prometheusURL != "" {
-			blockCheckers = append(blockCheckers, PrometheusBlockingChecker{promURL: prometheusURL, filter: alertFilter})
+			blockCheckers = append(blockCheckers, PrometheusBlockingChecker{promClient: promClient, filter: alertFilter})
 		}
 		if podSelectors != nil {
 			blockCheckers = append(blockCheckers, KubernetesBlockingChecker{client: client, nodename: nodeID, filter: podSelectors})
@@ -5,7 +5,10 @@ import (
 	"testing"

 	log "github.com/sirupsen/logrus"
+	"github.com/weaveworks/kured/pkg/alerts"
 	assert "gotest.tools/v3/assert"
+
+	papi "github.com/prometheus/client_golang/api"
 )

 type BlockingChecker struct {
@@ -23,7 +26,13 @@ func Test_rebootBlocked(t *testing.T) {
 	noCheckers := []RebootBlocker{}
 	nonblockingChecker := BlockingChecker{blocking: false}
 	blockingChecker := BlockingChecker{blocking: true}
-	brokenPrometheusClient := PrometheusBlockingChecker{promURL: "", filter: nil}
+
+	// Instantiate a prometheusClient with a broken_url
+	promClient, err := alerts.NewPromClient(papi.Config{Address: "broken_url"})
+	if err != nil {
+		log.Fatal("Can't create prometheusClient: ", err)
+	}
+	brokenPrometheusClient := PrometheusBlockingChecker{promClient: promClient, filter: nil}

 	type args struct {
 		blockers []RebootBlocker
@@ -9,6 +9,7 @@ require (
 	github.com/prometheus/common v0.29.0
 	github.com/sirupsen/logrus v1.8.1
 	github.com/spf13/cobra v1.2.1
+	github.com/stretchr/testify v1.6.1
 	golang.org/x/crypto v0.0.0-20210506145944-38f3c27a63bf // indirect
 	gotest.tools/v3 v3.0.3
 	k8s.io/api v0.20.5
@@ -7,22 +7,39 @@ import (
 	"sort"
 	"time"

-	"github.com/prometheus/client_golang/api"
+	papi "github.com/prometheus/client_golang/api"
 	v1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/common/model"
 )

-// PrometheusActiveAlerts returns a list of names of active (e.g. pending or firing) alerts, filtered
-// by the supplied regexp.
-func PrometheusActiveAlerts(prometheusURL string, filter *regexp.Regexp) ([]string, error) {
-	client, err := api.NewClient(api.Config{Address: prometheusURL})
+// PromClient is a wrapper around the Prometheus Client interface and implements the api
+// This way, the PromClient can be instantiated with the configuration the Client needs, and
+// the ability to use the methods the api has, like Query and so on.
+type PromClient struct {
+	papi papi.Client
+	api  v1.API
+}
+
+// NewPromClient creates a new client to the Prometheus API.
+// It returns an error on any problem.
+func NewPromClient(conf papi.Config) (*PromClient, error) {
+	promClient, err := papi.NewClient(conf)
 	if err != nil {
 		return nil, err
 	}
+	client := PromClient{papi: promClient, api: v1.NewAPI(promClient)}
+	return &client, nil
+}

-	queryAPI := v1.NewAPI(client)
+// ActiveAlerts is a method of type PromClient, it returns a list of names of active alerts
+// (e.g. pending or firing), filtered by the supplied regexp or by the includeLabels query.
+// filter by regexp means when the regex finds the alert-name; the alert is exluded from the
+// block-list and will NOT block rebooting. query by includeLabel means,
+// if the query finds an alert, it will include it to the block-list and it WILL block rebooting.
+func (p *PromClient) ActiveAlerts(filter *regexp.Regexp) ([]string, error) {

-	value, _, err := queryAPI.Query(context.Background(), "ALERTS", time.Now())
+	// get all alerts from prometheus
+	value, _, err := p.api.Query(context.Background(), "ALERTS", time.Now())
 	if err != nil {
 		return nil, err
 	}
@@ -42,7 +59,7 @@ func PrometheusActiveAlerts(prometheusURL string, filter *regexp.Regexp) ([]stri
 			for activeAlert := range activeAlertSet {
 				activeAlerts = append(activeAlerts, activeAlert)
 			}
-			sort.Sort(sort.StringSlice(activeAlerts))
+			sort.Strings(activeAlerts)

 			return activeAlerts, nil
 		}
@@ -0,0 +1,127 @@
+package alerts
+
+import (
+	"log"
+	"net/http"
+	"net/http/httptest"
+
+	"regexp"
+	"testing"
+
+	"github.com/prometheus/client_golang/api"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type MockResponse struct {
+	StatusCode int
+	Body       []byte
+}
+
+// MockServerProperties ties a mock response to a url and a method
+type MockServerProperties struct {
+	URI        string
+	HTTPMethod string
+	Response   MockResponse
+}
+
+// NewMockServer sets up a new MockServer with properties ad starts the server.
+func NewMockServer(props ...MockServerProperties) *httptest.Server {
+
+	handler := http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			for _, proc := range props {
+				_, err := w.Write(proc.Response.Body)
+				if err != nil {
+					log.Fatal(err)
+				}
+			}
+		})
+	return httptest.NewServer(handler)
+}
+
+func TestActiveAlerts(t *testing.T) {
+	responsebody := `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"ALERTS","alertname":"GatekeeperViolations","alertstate":"firing","severity":"warning","team":"platform-infra"},"value":[1622472933.973,"1"]},{"metric":{"__name__":"ALERTS","alertname":"PodCrashing-dev","alertstate":"firing","container":"deployment","instance":"1.2.3.4:8080","job":"kube-state-metrics","namespace":"dev","pod":"dev-deployment-78dcbmf25v","severity":"critical","team":"dev"},"value":[1622472933.973,"1"]},{"metric":{"__name__":"ALERTS","alertname":"PodRestart-dev","alertstate":"firing","container":"deployment","instance":"1.2.3.4:1234","job":"kube-state-metrics","namespace":"qa","pod":"qa-job-deployment-78dcbmf25v","severity":"warning","team":"qa"},"value":[1622472933.973,"1"]},{"metric":{"__name__":"ALERTS","alertname":"PrometheusTargetDown","alertstate":"firing","job":"kubernetes-pods","severity":"warning","team":"platform-infra"},"value":[1622472933.973,"1"]},{"metric":{"__name__":"ALERTS","alertname":"ScheduledRebootFailing","alertstate":"pending","severity":"warning","team":"platform-infra"},"value":[1622472933.973,"1"]}]}}`
+	addr := "http://localhost:10001"
+
+	for _, tc := range []struct {
+		it       string
+		rFilter  string
+		respBody string
+		aName    string
+		wantN    int
+	}{
+		{
+			it:       "should return no active alerts",
+			respBody: responsebody,
+			rFilter:  "",
+			wantN:    0,
+		},
+		{
+			it:       "should return a subset of all alerts",
+			respBody: responsebody,
+			rFilter:  "Pod",
+			wantN:    3,
+		},
+		{
+			it:       "should return all active alerts by regex",
+			respBody: responsebody,
+			rFilter:  "*",
+			wantN:    5,
+		},
+		{
+			it:       "should return all active alerts by regex filter",
+			respBody: responsebody,
+			rFilter:  "*",
+			wantN:    5,
+		},
+		{
+			it:       "should return ScheduledRebootFailing active alerts",
+			respBody: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"ALERTS","alertname":"ScheduledRebootFailing","alertstate":"pending","severity":"warning","team":"platform-infra"},"value":[1622472933.973,"1"]}]}}`,
+			aName:    "ScheduledRebootFailing",
+			rFilter:  "*",
+			wantN:    1,
+		},
+		{
+			it:       "should not return an active alert if RebootRequired is firing (regex filter)",
+			respBody: `{"status":"success","data":{"resultType":"vector","result":[{"metric":{"__name__":"ALERTS","alertname":"RebootRequired","alertstate":"pending","severity":"warning","team":"platform-infra"},"value":[1622472933.973,"1"]}]}}`,
+			rFilter:  "RebootRequired",
+			wantN:    0,
+		},
+	} {
+		// Start mockServer
+		mockServer := NewMockServer(MockServerProperties{
+			URI:        addr,
+			HTTPMethod: http.MethodPost,
+			Response: MockResponse{
+				Body: []byte(tc.respBody),
+			},
+		})
+		// Close mockServer after all connections are gone
+		defer mockServer.Close()
+
+		t.Run(tc.it, func(t *testing.T) {
+
+			// regex filter
+			regex, _ := regexp.Compile(tc.rFilter)
+
+			// instantiate the prometheus client with the mockserver-address
+			p, err := NewPromClient(api.Config{Address: mockServer.URL})
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			result, err := p.ActiveAlerts(regex)
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			// assert
+			assert.Equal(t, tc.wantN, len(result), "expected amount of alerts %v, got %v", tc.wantN, len(result))
+
+			if tc.aName != "" {
+				assert.Equal(t, tc.aName, result[0], "expected active alert %v, got %v", tc.aName, result[0])
+			}
+		})
+	}
+}