@@ -16,11 +16,16 @@ package framework
16
16
17
17
import (
18
18
"crypto/tls"
19
+ "fmt"
19
20
"io/ioutil"
20
21
"net/http"
22
+ "strconv"
21
23
"strings"
24
+ "testing"
25
+ "time"
22
26
23
27
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
28
+ "k8s.io/apimachinery/pkg/util/wait"
24
29
"k8s.io/client-go/kubernetes"
25
30
26
31
routev1 "github.com/openshift/client-go/route/clientset/versioned/typed/route/v1"
@@ -37,7 +42,7 @@ type PrometheusClient struct {
37
42
token string
38
43
}
39
44
40
- // NewPrometheusClient returns creates and returns a new PrometheusClient.
45
+ // NewPrometheusClient creates and returns a new PrometheusClient.
41
46
func NewPrometheusClient (
42
47
routeClient routev1.RouteV1Interface ,
43
48
kubeClient kubernetes.Interface ,
@@ -68,8 +73,9 @@ func NewPrometheusClient(
68
73
}, nil
69
74
}
70
75
71
- // Query makes a request against the Prometheus /api/v1/query endpoint.
72
- func (c * PrometheusClient ) Query (query string ) (int , error ) {
76
+ // Query runs an http get request against the Prometheus query api and returns
77
+ // the response body.
78
+ func (c * PrometheusClient ) Query (query string ) ([]byte , error ) {
73
79
tr := & http.Transport {
74
80
TLSClientConfig : & tls.Config {InsecureSkipVerify : true },
75
81
}
@@ -78,7 +84,7 @@ func (c *PrometheusClient) Query(query string) (int, error) {
78
84
79
85
req , err := http .NewRequest ("GET" , "https://" + c .host + "/api/v1/query" , nil )
80
86
if err != nil {
81
- return 0 , err
87
+ return nil , err
82
88
}
83
89
84
90
q := req .URL .Query ()
@@ -89,21 +95,103 @@ func (c *PrometheusClient) Query(query string) (int, error) {
89
95
90
96
resp , err := client .Do (req )
91
97
if err != nil {
92
- return 0 , err
98
+ return nil , err
93
99
}
94
100
95
101
defer resp .Body .Close ()
96
102
97
103
body , err := ioutil .ReadAll (resp .Body )
98
104
if err != nil {
99
- return 0 , err
105
+ return nil , err
100
106
}
101
107
108
+ return body , nil
109
+ }
110
+
111
+ // GetFirstValueFromPromQuery takes a query api response body and returns the
112
+ // value of the first timeseries. If body contains multiple timeseries
113
+ // GetFirstValueFromPromQuery errors.
114
+ func GetFirstValueFromPromQuery (body []byte ) (int , error ) {
102
115
res , err := gabs .ParseJSON (body )
103
116
if err != nil {
104
117
return 0 , err
105
118
}
106
119
107
- n , err := res .ArrayCountP ("data.result" )
108
- return n , err
120
+ count , err := res .ArrayCountP ("data.result" )
121
+ if err != nil {
122
+ return 0 , err
123
+ }
124
+
125
+ if count != 1 {
126
+ return 0 , fmt .Errorf ("expected body to contain single timeseries but got %v" , count )
127
+ }
128
+
129
+ timeseries , err := res .ArrayElementP (0 , "data.result" )
130
+ if err != nil {
131
+ return 0 , err
132
+ }
133
+
134
+ value , err := timeseries .ArrayElementP (1 , "value" )
135
+ if err != nil {
136
+ return 0 , err
137
+ }
138
+
139
+ v , err := strconv .Atoi (value .Data ().(string ))
140
+ if err != nil {
141
+ return 0 , fmt .Errorf ("failed to parse query value: %v" , err )
142
+ }
143
+
144
+ return v , nil
145
+ }
146
+
147
+ // WaitForQueryReturnGreaterEqualOne see WaitForQueryReturn.
148
+ func (c * PrometheusClient ) WaitForQueryReturnGreaterEqualOne (t * testing.T , timeout time.Duration , query string ) {
149
+ c .WaitForQueryReturn (t , timeout , query , func (v int ) error {
150
+ if v >= 1 {
151
+ return nil
152
+ }
153
+
154
+ return fmt .Errorf ("expected value to equal or greater than 1 but got %v" , v )
155
+ })
156
+ }
157
+
158
+ // WaitForQueryReturnOne see WaitForQueryReturn.
159
+ func (c * PrometheusClient ) WaitForQueryReturnOne (t * testing.T , timeout time.Duration , query string ) {
160
+ c .WaitForQueryReturn (t , timeout , query , func (v int ) error {
161
+ if v == 1 {
162
+ return nil
163
+ }
164
+
165
+ return fmt .Errorf ("expected value to equal 1 but got %v" , v )
166
+ })
167
+ }
168
+
169
+ // WaitForQueryReturn waits for a given PromQL query for a given time interval
170
+ // and validates the **first and only** result with the given validate function.
171
+ func (c * PrometheusClient ) WaitForQueryReturn (t * testing.T , timeout time.Duration , query string , validate func (int ) error ) {
172
+ err := wait .Poll (5 * time .Second , timeout , func () (bool , error ) {
173
+ defer t .Log ("---------------------------\n " )
174
+ body , err := c .Query (query )
175
+ if err != nil {
176
+ return false , err
177
+ }
178
+
179
+ v , err := GetFirstValueFromPromQuery (body )
180
+ if err != nil {
181
+ t .Logf ("failed to extract first value from query response for query %q: %v" , query , err )
182
+ return false , nil
183
+ }
184
+
185
+ if err := validate (v ); err != nil {
186
+ t .Logf ("unexpected value for query %q: %v" , query , err )
187
+ return false , nil
188
+ }
189
+
190
+ t .Logf ("query %q succeeded" , query )
191
+ return true , nil
192
+ })
193
+
194
+ if err != nil {
195
+ t .Fatal (err )
196
+ }
109
197
}
0 commit comments