11package prom
22
33import (
4+ "reflect"
45 "testing"
56 "time"
67
78 "github.com/prometheus/client_golang/prometheus"
89 dto "github.com/prometheus/client_model/go"
910 "github.com/stretchr/testify/assert"
10- "github.com/stretchr/testify/require"
1111)
1212
1313var s = NewPrometheusSink ()
@@ -16,21 +16,21 @@ func TestFlushCounter(t *testing.T) {
1616 s .FlushCounter ("ratelimit_server.ShouldRateLimit.total_requests" , 1 )
1717 assert .Eventually (t , func () bool {
1818 metricFamilies , err := prometheus .DefaultGatherer .Gather ()
19- require .NoError (t , err )
19+ if err != nil {
20+ return false
21+ }
2022
2123 metrics := make (map [string ]* dto.MetricFamily )
2224 for _ , metricFamily := range metricFamilies {
2325 metrics [* metricFamily .Name ] = metricFamily
2426 }
2527
2628 m , ok := metrics ["ratelimit_service_total_requests" ]
27- require .True (t , ok )
28- require .Len (t , m .Metric , 1 )
29- require .Equal (t , map [string ]string {
30- "grpc_method" : "ShouldRateLimit" ,
31- }, toMap (m .Metric [0 ].Label ))
32- require .Equal (t , 1.0 , * m .Metric [0 ].Counter .Value )
33- return true
29+ if ! ok || len (m .Metric ) != 1 {
30+ return false
31+ }
32+ return toMap (m .Metric [0 ].Label )["grpc_method" ] == "ShouldRateLimit" &&
33+ * m .Metric [0 ].Counter .Value == 1.0
3434 }, time .Second , time .Millisecond )
3535}
3636
@@ -49,71 +49,75 @@ func TestFlushCounterWithDifferentLabels(t *testing.T) {
4949 s .FlushCounter ("ratelimit.service.rate_limit.domain1.key3_val3.key4_val4.key5_val5.over_limit" , 2 )
5050 assert .Eventually (t , func () bool {
5151 metricFamilies , err := prometheus .DefaultGatherer .Gather ()
52- require .NoError (t , err )
52+ if err != nil {
53+ return false
54+ }
5355
5456 metrics := make (map [string ]* dto.MetricFamily )
5557 for _ , metricFamily := range metricFamilies {
5658 metrics [* metricFamily .Name ] = metricFamily
5759 }
5860
5961 m , ok := metrics ["ratelimit_service_rate_limit_over_limit" ]
60- require . True ( t , ok )
61- require . Len ( t , m . Metric , 3 )
62- require . Equal ( t , 1.0 , * m . Metric [ 0 ]. Counter . Value )
63- require . Equal ( t , map [ string ] string {
64- "domain" : "domain1" ,
65- "key1 " "key1_val1 "
66- }, toMap ( m . Metric [ 0 ]. Label ))
67- require . Equal ( t , 2.0 , * m . Metric [ 1 ]. Counter . Value )
68- require . Equal ( t , map [ string ] string {
69- "domain" : "domain1" ,
70- "key1 " "key1_val1 "
71- "key2 ""key2_val2 " ,
72- }, toMap ( m . Metric [ 1 ]. Label ))
73- require . Equal ( t , 3.0 , * m . Metric [ 2 ]. Counter . Value )
74- require . Equal ( t , map [ string ] string {
75- "domain" : "domain1" ,
76- "key1 " "key3_val3 "
77- "key2 ""key4_val4 " ,
78- }, toMap ( m . Metric [ 2 ]. Label ))
79- return true
62+ if ! ok || len ( m . Metric ) != 3 {
63+ return false
64+ }
65+ return * m . Metric [ 0 ]. Counter . Value == 1.0 &&
66+ reflect . DeepEqual ( toMap ( m . Metric [ 0 ]. Label ), map [ string ] string {
67+ "domain ""domain1 "
68+ "key1" : "key1_val1" ,
69+ }) &&
70+ * m . Metric [ 1 ]. Counter . Value == 2.0 &&
71+ reflect . DeepEqual ( toMap ( m . Metric [ 1 ]. Label ), map [ string ] string {
72+ "domain ""domain1 "
73+ "key1 ""key1_val1 " ,
74+ "key2" : "key2_val2" ,
75+ }) &&
76+ * m . Metric [ 2 ]. Counter . Value == 3.0 &&
77+ reflect . DeepEqual ( toMap ( m . Metric [ 2 ]. Label ), map [ string ] string {
78+ "domain ""domain1 "
79+ "key1 ""key3_val3 " ,
80+ "key2" : "key4_val4" ,
81+ })
8082 }, time .Second , time .Millisecond )
8183}
8284
8385func TestFlushGauge (t * testing.T ) {
8486 s .FlushGauge ("ratelimit.service.rate_limit.domain1.key1.test_gauge" , 1 )
8587 metricFamilies , err := prometheus .DefaultGatherer .Gather ()
86- require .NoError (t , err )
88+ assert .NoError (t , err )
8789
8890 metrics := make (map [string ]* dto.MetricFamily )
8991 for _ , metricFamily := range metricFamilies {
9092 metrics [* metricFamily .Name ] = metricFamily
9193 }
9294
9395 _ , ok := metrics ["ratelimit_service_rate_limit_test_gauge" ]
94- require .False (t , ok )
96+ assert .False (t , ok )
9597}
9698
9799func TestFlushTimer (t * testing.T ) {
98100 s .FlushTimer ("ratelimit.service.rate_limit.mongo_cps.database_users.total_hits" , 1 )
99101 assert .Eventually (t , func () bool {
100102 metricFamilies , err := prometheus .DefaultGatherer .Gather ()
101- require .NoError (t , err )
103+ if err != nil {
104+ return false
105+ }
102106
103107 metrics := make (map [string ]* dto.MetricFamily )
104108 for _ , metricFamily := range metricFamilies {
105109 metrics [* metricFamily .Name ] = metricFamily
106110 }
107111
108112 m , ok := metrics ["ratelimit_service_rate_limit_total_hits" ]
109- require . True ( t , ok )
110- require . Len ( t , m . Metric , 1 )
111- require . Equal ( t , uint64 ( 1 ), * m . Metric [ 0 ]. Histogram . SampleCount )
112- require . Equal ( t , map [ string ] string {
113- "domain" : "mongo_cps" ,
114- "key1 " "database_users "
115- }, toMap ( m . Metric [ 0 ]. Label ))
116- require . Equal ( t , 1.0 , * m . Metric [ 0 ]. Histogram . SampleSum )
117- return true
113+ if ! ok || len ( m . Metric ) != 1 {
114+ return false
115+ }
116+ return * m . Metric [ 0 ]. Histogram . SampleCount == uint64 ( 1 ) &&
117+ reflect . DeepEqual ( toMap ( m . Metric [ 0 ]. Label ), map [ string ] string {
118+ "domain ""mongo_cps "
119+ "key1" : "database_users" ,
120+ }) &&
121+ * m . Metric [ 0 ]. Histogram . SampleSum == 1.0
118122 }, time .Second , time .Millisecond )
119123}
0 commit comments