Skip to content

Commit 5270be9

Browse files
yulong-dbyulong-db
committed
add unit tests
1 parent 06f1420 commit 5270be9

File tree

3 files changed

+163
-11
lines changed

3 files changed

+163
-11
lines changed

‎pkg/query/aggregation_label_rewriter.go‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import (
1313

1414
const (
1515
aggregationLabelName = "__rollup__"
16-
aggregatedMetricNameGlobPattern = "*:(aggr|sum|count|avg|min|max)"
16+
aggregatedMetricNameGlobPattern = "{*:aggr,*:sum,*:count,*:avg,*:min,*:max}"
1717
)
1818

1919
type AggregationLabelRewriter struct {
@@ -60,6 +60,9 @@ func newAggregationLabelRewriterMetrics(reg prometheus.Registerer, desiredLabelV
6060
}
6161

6262
func NewAggregationLabelRewriter(logger log.Logger, reg prometheus.Registerer, desiredLabelValue string) *AggregationLabelRewriter {
63+
if logger == nil {
64+
logger = log.NewNopLogger()
65+
}
6366
g := glob.MustCompile(aggregatedMetricNameGlobPattern)
6467
return &AggregationLabelRewriter{
6568
enabled: desiredLabelValue != "",
@@ -76,6 +79,7 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
7679
needsRewrite := false
7780
skipReason := "no-name-matcher"
7881
var aggregationLabelMatcher *labels.Matcher
82+
aggregationLabelIndex := -1
7983
for i := 0; i < len(ms); i++ {
8084
m := ms[i]
8185
if !needsRewrite && m.Name == labels.MetricName {
@@ -94,20 +98,20 @@ func (a *AggregationLabelRewriter) Rewrite(ms []*labels.Matcher) []*labels.Match
9498
break
9599
} else if m.Name == aggregationLabelName {
96100
aggregationLabelMatcher = m
101+
aggregationLabelIndex = i
97102
}
98103
}
99104
if needsRewrite {
105+
newMatcher := &labels.Matcher{
106+
Name: aggregationLabelName,
107+
Type: labels.MatchEqual,
108+
Value: a.desiredLabelValue,
109+
}
100110
if aggregationLabelMatcher != nil {
101111
a.metrics.aggregationLabelRewrittenCount.WithLabelValues(aggregationLabelMatcher.Value).Inc()
102-
aggregationLabelMatcher.Type = labels.MatchEqual
103-
aggregationLabelMatcher.Value = a.desiredLabelValue
112+
ms[aggregationLabelIndex] = newMatcher
104113
} else {
105114
a.metrics.aggregationLabelAddedCount.Inc()
106-
newMatcher := &labels.Matcher{
107-
Name: aggregationLabelName,
108-
Type: labels.MatchEqual,
109-
Value: a.desiredLabelValue,
110-
}
111115
ms = append(ms, newMatcher)
112116
}
113117
a.metrics.aggregationLabelRewriterRuntimeSeconds.WithLabelValues("true").Observe(

‎pkg/query/aggregation_label_rewriter_test.go‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
package query
2+
3+
import (
4+
"testing"
5+
6+
"github.com/efficientgo/core/testutil"
7+
"github.com/prometheus/client_golang/prometheus"
8+
"github.com/prometheus/prometheus/model/labels"
9+
)
10+
11+
func TestAggregationLabelRewriter_Rewrite(t *testing.T) {
12+
t.Parallel()
13+
14+
for _, tc := range []struct {
15+
name string
16+
desiredLabelValue string // Empty means disabled
17+
inputMatchers []*labels.Matcher
18+
expectedMatchers []*labels.Matcher
19+
expectedSkipCount float64
20+
expectedAddCount float64
21+
expectedRewriteMap map[string]float64
22+
}{
23+
{
24+
name: "disabled rewriter should not modify label matchers",
25+
desiredLabelValue: "",
26+
inputMatchers: []*labels.Matcher{
27+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
28+
},
29+
expectedMatchers: []*labels.Matcher{
30+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
31+
},
32+
},
33+
{
34+
name: "should add label for aggregated metric if no existing aggregation label",
35+
desiredLabelValue: "5m",
36+
inputMatchers: []*labels.Matcher{
37+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
38+
},
39+
expectedMatchers: []*labels.Matcher{
40+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
41+
labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "5m"),
42+
},
43+
expectedAddCount: 1,
44+
},
45+
{
46+
name: "should rewrite existing aggregation label for aggregated metric",
47+
desiredLabelValue: "5m",
48+
inputMatchers: []*labels.Matcher{
49+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
50+
labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "1h"),
51+
},
52+
expectedMatchers: []*labels.Matcher{
53+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test:sum"),
54+
labels.MustNewMatcher(labels.MatchEqual, "__rollup__", "5m"),
55+
},
56+
expectedRewriteMap: map[string]float64{"1h": 1},
57+
},
58+
{
59+
name: "should skip non-aggregated metric",
60+
desiredLabelValue: "5m",
61+
inputMatchers: []*labels.Matcher{
62+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test_metric"),
63+
},
64+
expectedMatchers: []*labels.Matcher{
65+
labels.MustNewMatcher(labels.MatchEqual, "__name__", "test_metric"),
66+
},
67+
expectedSkipCount: 1,
68+
},
69+
{
70+
name: "should skip non-equal name matcher",
71+
desiredLabelValue: "5m",
72+
inputMatchers: []*labels.Matcher{
73+
labels.MustNewMatcher(labels.MatchRegexp, "__name__", "test:sum"),
74+
},
75+
expectedMatchers: []*labels.Matcher{
76+
labels.MustNewMatcher(labels.MatchRegexp, "__name__", "test:sum"),
77+
},
78+
expectedSkipCount: 1,
79+
},
80+
{
81+
name: "should skip when no name matcher",
82+
desiredLabelValue: "5m",
83+
inputMatchers: []*labels.Matcher{
84+
labels.MustNewMatcher(labels.MatchEqual, "job", "prometheus"),
85+
},
86+
expectedMatchers: []*labels.Matcher{
87+
labels.MustNewMatcher(labels.MatchEqual, "job", "prometheus"),
88+
},
89+
expectedSkipCount: 1,
90+
},
91+
} {
92+
t.Run(tc.name, func(t *testing.T) {
93+
reg := prometheus.NewRegistry()
94+
rewriter := NewAggregationLabelRewriter(
95+
nil,
96+
reg,
97+
tc.desiredLabelValue,
98+
)
99+
100+
result := rewriter.Rewrite(tc.inputMatchers)
101+
testutil.Equals(t, len(tc.expectedMatchers), len(result))
102+
for i := range tc.inputMatchers {
103+
testutil.Equals(t, tc.expectedMatchers[i].Name, result[i].Name)
104+
testutil.Equals(t, tc.expectedMatchers[i].Type, result[i].Type)
105+
testutil.Equals(t, tc.expectedMatchers[i].Value, result[i].Value)
106+
}
107+
108+
metrics, err := reg.Gather()
109+
testutil.Ok(t, err)
110+
111+
if tc.expectedSkipCount > 0 {
112+
var skipCount float64
113+
for _, m := range metrics {
114+
if m.GetName() == "skipped_total" {
115+
skipCount += *m.Metric[0].Counter.Value
116+
}
117+
}
118+
testutil.Equals(t, tc.expectedSkipCount, skipCount)
119+
}
120+
121+
if tc.expectedAddCount > 0 {
122+
var addCount float64
123+
for _, m := range metrics {
124+
if m.GetName() == "label_added_total" {
125+
addCount += *m.Metric[0].Counter.Value
126+
}
127+
}
128+
testutil.Equals(t, tc.expectedAddCount, addCount)
129+
}
130+
131+
if len(tc.expectedRewriteMap) > 0 {
132+
for _, m := range metrics {
133+
if m.GetName() == "label_rewritten_total" {
134+
for _, metric := range m.Metric {
135+
oldValue := ""
136+
for _, label := range metric.Label {
137+
if *label.Name == "old_value" {
138+
oldValue = *label.Value
139+
break
140+
}
141+
}
142+
if count, ok := tc.expectedRewriteMap[oldValue]; ok {
143+
testutil.Equals(t, count, *metric.Counter.Value)
144+
}
145+
}
146+
}
147+
}
148+
}
149+
})
150+
}
151+
}

‎pkg/query/querier.go‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -211,9 +211,6 @@ func newQuerierWithOpts(
211211
if logger == nil {
212212
logger = log.NewNopLogger()
213213
}
214-
if reg == nil {
215-
reg = prometheus.NewRegistry()
216-
}
217214
rl := make(map[string]struct{})
218215
for _, replicaLabel := range replicaLabels {
219216
rl[replicaLabel] = struct{}{}

0 commit comments

Comments
 (0)