ratelimit/src/stats/prom/prometheus_sink_test.go at 992569f4575821134f17e472fe502d3d145447a5 · Revolyssup/ratelimit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
package prom

import (
	"reflect"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/assert"
)

var s = NewPrometheusSink()

func TestFlushCounter(t *testing.T) {
	s.FlushCounter("ratelimit_server.ShouldRateLimit.total_requests", 1)
	assert.Eventually(t, func() bool {
		metricFamilies, err := prometheus.DefaultGatherer.Gather()
		if err != nil {
			return false
		}

		metrics := make(map[string]*dto.MetricFamily)
		for _, metricFamily := range metricFamilies {
			metrics[*metricFamily.Name] = metricFamily
		}

		m, ok := metrics["ratelimit_service_total_requests"]
		if !ok || len(m.Metric) != 1 {
			return false
		}
		return toMap(m.Metric[0].Label)["grpc_method"] == "ShouldRateLimit" &&
			*m.Metric[0].Counter.Value == 1.0
	}, time.Second, time.Millisecond)
}

func toMap(labels []*dto.LabelPair) map[string]string {
	m := make(map[string]string)
	for _, l := range labels {
		m[*l.Name] = *l.Value
	}
	return m
}

func TestFlushCounterWithDifferentLabels(t *testing.T) {
	s.FlushCounter("ratelimit.service.rate_limit.domain1.key1_val1.over_limit", 1)
	s.FlushCounter("ratelimit.service.rate_limit.domain1.key1_val1.key2_val2.over_limit", 2)
	s.FlushCounter("ratelimit.service.rate_limit.domain1.key3_val3.key4_val4.over_limit", 1)
	s.FlushCounter("ratelimit.service.rate_limit.domain1.key3_val3.key4_val4.key5_val5.over_limit", 2)
	assert.Eventually(t, func() bool {
		metricFamilies, err := prometheus.DefaultGatherer.Gather()
		if err != nil {
			return false
		}

		metrics := make(map[string]*dto.MetricFamily)
		for _, metricFamily := range metricFamilies {
			metrics[*metricFamily.Name] = metricFamily
		}

		m, ok := metrics["ratelimit_service_rate_limit_over_limit"]
		if !ok || len(m.Metric) != 3 {
			return false
		}
		return *m.Metric[0].Counter.Value == 1.0 &&
			reflect.DeepEqual(toMap(m.Metric[0].Label), map[string]string{
				"domain": "domain1",
				"key1":   "key1_val1",
			}) &&
			*m.Metric[1].Counter.Value == 2.0 &&
			reflect.DeepEqual(toMap(m.Metric[1].Label), map[string]string{
				"domain": "domain1",
				"key1":   "key1_val1",
				"key2":   "key2_val2",
			}) &&
			*m.Metric[2].Counter.Value == 3.0 &&
			reflect.DeepEqual(toMap(m.Metric[2].Label), map[string]string{
				"domain": "domain1",
				"key1":   "key3_val3",
				"key2":   "key4_val4",
			})
	}, time.Second, time.Millisecond)
}

func TestFlushGauge(t *testing.T) {
	s.FlushGauge("ratelimit.service.rate_limit.domain1.key1.test_gauge", 1)
	metricFamilies, err := prometheus.DefaultGatherer.Gather()
	assert.NoError(t, err)

	metrics := make(map[string]*dto.MetricFamily)
	for _, metricFamily := range metricFamilies {
		metrics[*metricFamily.Name] = metricFamily
	}

	_, ok := metrics["ratelimit_service_rate_limit_test_gauge"]
	assert.False(t, ok)
}

func TestFlushTimer(t *testing.T) {
	s.FlushTimer("ratelimit.service.rate_limit.mongo_cps.database_users.total_hits", 1)
	assert.Eventually(t, func() bool {
		metricFamilies, err := prometheus.DefaultGatherer.Gather()
		if err != nil {
			return false
		}

		metrics := make(map[string]*dto.MetricFamily)
		for _, metricFamily := range metricFamilies {
			metrics[*metricFamily.Name] = metricFamily
		}

		m, ok := metrics["ratelimit_service_rate_limit_total_hits"]
		if !ok || len(m.Metric) != 1 {
			return false
		}
		return *m.Metric[0].Histogram.SampleCount == uint64(1) &&
			reflect.DeepEqual(toMap(m.Metric[0].Label), map[string]string{
				"domain": "mongo_cps",
				"key1":   "database_users",
			}) &&
			*m.Metric[0].Histogram.SampleSum == 1.0
	}, time.Second, time.Millisecond)
}