github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/querier/queryrange/queryable_test.go (about) 1 package queryrange 2 3 import ( 4 "context" 5 "testing" 6 7 "github.com/pkg/errors" 8 "github.com/prometheus/prometheus/pkg/labels" 9 "github.com/prometheus/prometheus/promql/parser" 10 "github.com/stretchr/testify/require" 11 12 "github.com/cortexproject/cortex/pkg/cortexpb" 13 "github.com/cortexproject/cortex/pkg/querier/astmapper" 14 ) 15 16 func TestSelect(t *testing.T) { 17 var testExpr = []struct { 18 name string 19 querier *ShardedQuerier 20 fn func(*testing.T, *ShardedQuerier) 21 }{ 22 { 23 name: "errors non embedded query", 24 querier: mkQuerier( 25 nil, 26 ), 27 fn: func(t *testing.T, q *ShardedQuerier) { 28 set := q.Select(false, nil) 29 require.EqualError(t, set.Err(), nonEmbeddedErrMsg) 30 }, 31 }, 32 { 33 name: "replaces query", 34 querier: mkQuerier(mockHandlerWith( 35 &PrometheusResponse{}, 36 nil, 37 )), 38 fn: func(t *testing.T, q *ShardedQuerier) { 39 40 expected := &PrometheusResponse{ 41 Status: "success", 42 Data: PrometheusData{ 43 ResultType: string(parser.ValueTypeVector), 44 }, 45 } 46 47 // override handler func to assert new query has been substituted 48 q.Handler = HandlerFunc( 49 func(ctx context.Context, req Request) (Response, error) { 50 require.Equal(t, `http_requests_total{cluster="prod"}`, req.GetQuery()) 51 return expected, nil 52 }, 53 ) 54 55 encoded, err := astmapper.JSONCodec.Encode([]string{`http_requests_total{cluster="prod"}`}) 56 require.Nil(t, err) 57 set := q.Select( 58 false, 59 nil, 60 exactMatch("__name__", astmapper.EmbeddedQueriesMetricName), 61 exactMatch(astmapper.QueryLabel, encoded), 62 ) 63 require.Nil(t, set.Err()) 64 }, 65 }, 66 { 67 name: "propagates response error", 68 querier: mkQuerier(mockHandlerWith( 69 &PrometheusResponse{ 70 Error: "SomeErr", 71 }, 72 nil, 73 )), 74 fn: func(t *testing.T, q *ShardedQuerier) { 75 encoded, err := astmapper.JSONCodec.Encode([]string{`http_requests_total{cluster="prod"}`}) 76 require.Nil(t, err) 77 set := q.Select( 78 false, 79 nil, 80 exactMatch("__name__", astmapper.EmbeddedQueriesMetricName), 81 exactMatch(astmapper.QueryLabel, encoded), 82 ) 83 require.EqualError(t, set.Err(), "SomeErr") 84 }, 85 }, 86 { 87 name: "returns SeriesSet", 88 querier: mkQuerier(mockHandlerWith( 89 &PrometheusResponse{ 90 Data: PrometheusData{ 91 ResultType: string(parser.ValueTypeVector), 92 Result: []SampleStream{ 93 { 94 Labels: []cortexpb.LabelAdapter{ 95 {Name: "a", Value: "a1"}, 96 {Name: "b", Value: "b1"}, 97 }, 98 Samples: []cortexpb.Sample{ 99 { 100 Value: 1, 101 TimestampMs: 1, 102 }, 103 { 104 Value: 2, 105 TimestampMs: 2, 106 }, 107 }, 108 }, 109 { 110 Labels: []cortexpb.LabelAdapter{ 111 {Name: "a", Value: "a1"}, 112 {Name: "b", Value: "b1"}, 113 }, 114 Samples: []cortexpb.Sample{ 115 { 116 Value: 8, 117 TimestampMs: 1, 118 }, 119 { 120 Value: 9, 121 TimestampMs: 2, 122 }, 123 }, 124 }, 125 }, 126 }, 127 }, 128 nil, 129 )), 130 fn: func(t *testing.T, q *ShardedQuerier) { 131 encoded, err := astmapper.JSONCodec.Encode([]string{`http_requests_total{cluster="prod"}`}) 132 require.Nil(t, err) 133 set := q.Select( 134 false, 135 nil, 136 exactMatch("__name__", astmapper.EmbeddedQueriesMetricName), 137 exactMatch(astmapper.QueryLabel, encoded), 138 ) 139 require.Nil(t, set.Err()) 140 require.Equal( 141 t, 142 NewSeriesSet([]SampleStream{ 143 { 144 Labels: []cortexpb.LabelAdapter{ 145 {Name: "a", Value: "a1"}, 146 {Name: "b", Value: "b1"}, 147 }, 148 Samples: []cortexpb.Sample{ 149 { 150 Value: 1, 151 TimestampMs: 1, 152 }, 153 { 154 Value: 2, 155 TimestampMs: 2, 156 }, 157 }, 158 }, 159 { 160 Labels: []cortexpb.LabelAdapter{ 161 {Name: "a", Value: "a1"}, 162 {Name: "b", Value: "b1"}, 163 }, 164 Samples: []cortexpb.Sample{ 165 { 166 Value: 8, 167 TimestampMs: 1, 168 }, 169 { 170 Value: 9, 171 TimestampMs: 2, 172 }, 173 }, 174 }, 175 }), 176 set, 177 ) 178 }, 179 }, 180 } 181 182 for _, c := range testExpr { 183 t.Run(c.name, func(t *testing.T) { 184 c.fn(t, c.querier) 185 }) 186 } 187 } 188 189 func TestSelectConcurrent(t *testing.T) { 190 for _, c := range []struct { 191 name string 192 queries []string 193 err error 194 }{ 195 { 196 name: "concats queries", 197 queries: []string{ 198 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="0_of_3",baz="blip"}[1m]))`, 199 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="1_of_3",baz="blip"}[1m]))`, 200 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="2_of_3",baz="blip"}[1m]))`, 201 }, 202 err: nil, 203 }, 204 { 205 name: "errors", 206 queries: []string{ 207 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="0_of_3",baz="blip"}[1m]))`, 208 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="1_of_3",baz="blip"}[1m]))`, 209 `sum by(__cortex_shard__) (rate(bar1{__cortex_shard__="2_of_3",baz="blip"}[1m]))`, 210 }, 211 err: errors.Errorf("some-err"), 212 }, 213 } { 214 215 t.Run(c.name, func(t *testing.T) { 216 // each request will return a single samplestream 217 querier := mkQuerier(mockHandlerWith(&PrometheusResponse{ 218 Data: PrometheusData{ 219 ResultType: string(parser.ValueTypeVector), 220 Result: []SampleStream{ 221 { 222 Labels: []cortexpb.LabelAdapter{ 223 {Name: "a", Value: "1"}, 224 }, 225 Samples: []cortexpb.Sample{ 226 { 227 Value: 1, 228 TimestampMs: 1, 229 }, 230 }, 231 }, 232 }, 233 }, 234 }, c.err)) 235 236 encoded, err := astmapper.JSONCodec.Encode(c.queries) 237 require.Nil(t, err) 238 set := querier.Select( 239 false, 240 nil, 241 exactMatch("__name__", astmapper.EmbeddedQueriesMetricName), 242 exactMatch(astmapper.QueryLabel, encoded), 243 ) 244 245 if c.err != nil { 246 require.EqualError(t, set.Err(), c.err.Error()) 247 return 248 } 249 250 var ct int 251 for set.Next() { 252 ct++ 253 } 254 require.Equal(t, len(c.queries), ct) 255 }) 256 } 257 } 258 259 func exactMatch(k, v string) *labels.Matcher { 260 m, err := labels.NewMatcher(labels.MatchEqual, k, v) 261 if err != nil { 262 panic(err) 263 } 264 return m 265 266 } 267 268 func mkQuerier(handler Handler) *ShardedQuerier { 269 return &ShardedQuerier{Ctx: context.Background(), Req: &PrometheusRequest{}, Handler: handler, ResponseHeaders: map[string][]string{}} 270 }