Antonio Mika
·
2026-10-25
1package router
2
3import (
4 "net/http"
5 "net/http/httptest"
6 "strings"
7 "testing"
8
9 "github.com/prometheus/client_golang/prometheus"
10)
11
12// metricValue sums the counters, or histogram counts, in reg that match
13// labels.
14func metricValue(t *testing.T, reg *prometheus.Registry, name string, labels map[string]string) float64 {
15 t.Helper()
16 families, err := reg.Gather()
17 if err != nil {
18 t.Fatalf("gather: %v", err)
19 }
20 var total float64
21 for _, family := range families {
22 if family.GetName() != name {
23 continue
24 }
25 metrics:
26 for _, m := range family.GetMetric() {
27 for _, l := range m.GetLabel() {
28 if want, ok := labels[l.GetName()]; ok && want != l.GetValue() {
29 continue metrics
30 }
31 }
32 if h := m.GetHistogram(); h != nil {
33 total += float64(h.GetSampleCount())
34 } else {
35 total += m.GetCounter().GetValue()
36 }
37 }
38 }
39 return total
40}
41
42// hostOwner reads "user-project.example.com" the way pgs reads its
43// subdomains.
44func hostOwner(r *http.Request) (string, string) {
45 user, resource, _ := strings.Cut(strings.Split(r.Host, ".")[0], "-")
46 return user, resource
47}
48
49func serve(h http.Handler, method, target string) *httptest.ResponseRecorder {
50 rec := httptest.NewRecorder()
51 h.ServeHTTP(rec, httptest.NewRequest(method, target, nil))
52 return rec
53}
54
55func TestUserMetricsSetUser(t *testing.T) {
56 reg := prometheus.NewRegistry()
57 m := NewUserMetrics("test", reg, UserMetricsOpts{})
58 h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
59 SetUser(r, "alice")
60 w.WriteHeader(http.StatusCreated)
61 _, _ = w.Write([]byte("hello"))
62 }))
63
64 serve(h, http.MethodPost, "/")
65
66 labels := map[string]string{"app": "test", "user": "alice", "status": "201", "method": "POST"}
67 if got := metricValue(t, reg, "pico_http_requests_total", labels); got != 1 {
68 t.Errorf("requests = %v, want 1", got)
69 }
70 if got := metricValue(t, reg, "pico_http_request_duration_seconds", labels); got != 1 {
71 t.Errorf("duration samples = %v, want 1", got)
72 }
73 if got := metricValue(t, reg, "pico_http_response_bytes_total", map[string]string{"user": "alice"}); got != 5 {
74 t.Errorf("bytes = %v, want 5", got)
75 }
76}
77
78func TestUserMetricsSiteUser(t *testing.T) {
79 reg := prometheus.NewRegistry()
80 m := NewUserMetrics("test", reg, UserMetricsOpts{Owner: hostOwner})
81 // The handler confirms alice on her first request; later ones stand in
82 // for cache hits that never reach it.
83 confirmed := false
84 h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
85 if !confirmed && r.Host == "alice.example.com" {
86 SetUser(r, "alice")
87 confirmed = true
88 }
89 }))
90
91 serve(h, http.MethodGet, "http://ghost.example.com/")
92 serve(h, http.MethodGet, "http://alice.example.com/")
93 serve(h, http.MethodGet, "http://alice.example.com/")
94
95 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice"}); got != 2 {
96 t.Errorf("alice requests = %v, want 2", got)
97 }
98 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": ""}); got != 1 {
99 t.Errorf("unlabeled requests = %v, want 1", got)
100 }
101 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "ghost"}); got != 0 {
102 t.Errorf("ghost requests = %v, want 0", got)
103 }
104}
105
106func TestUserMetricsTrustSuccess(t *testing.T) {
107 reg := prometheus.NewRegistry()
108 m := NewUserMetrics("test", reg, UserMetricsOpts{
109 Owner: hostOwner,
110 TrustSuccess: true,
111 })
112 h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
113 if r.Host == "ghost.example.com" {
114 http.NotFound(w, r)
115 }
116 }))
117
118 serve(h, http.MethodGet, "http://bob-docs.example.com/")
119 serve(h, http.MethodGet, "http://ghost.example.com/")
120
121 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "bob", "status": "200"}); got != 1 {
122 t.Errorf("bob requests = %v, want 1", got)
123 }
124 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "", "status": "404"}); got != 1 {
125 t.Errorf("unlabeled 404s = %v, want 1", got)
126 }
127 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "bob", "resource": "docs"}); got != 1 {
128 t.Errorf("bob/docs requests = %v, want 1", got)
129 }
130}
131
132func TestUserMetricsSetResource(t *testing.T) {
133 reg := prometheus.NewRegistry()
134 m := NewUserMetrics("test", reg, UserMetricsOpts{})
135 h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
136 SetUser(r, "alice")
137 if r.URL.Path == "/blog" {
138 SetResource(r, "blog")
139 }
140 _, _ = w.Write([]byte("hello"))
141 }))
142
143 serve(h, http.MethodGet, "/blog")
144 serve(h, http.MethodGet, "/")
145
146 labels := map[string]string{"app": "test", "user": "alice", "resource": "blog"}
147 if got := metricValue(t, reg, "pico_http_requests_total", labels); got != 1 {
148 t.Errorf("alice/blog requests = %v, want 1", got)
149 }
150 if got := metricValue(t, reg, "pico_http_response_bytes_total", labels); got != 5 {
151 t.Errorf("alice/blog bytes = %v, want 5", got)
152 }
153 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice"}); got != 2 {
154 t.Errorf("alice requests = %v, want 2", got)
155 }
156}
157
158func TestUserMetricsResourceFromHost(t *testing.T) {
159 reg := prometheus.NewRegistry()
160 m := NewUserMetrics("test", reg, UserMetricsOpts{Owner: hostOwner})
161 // The handler confirms alice's blog on its first request; later ones
162 // stand in for cache hits that never reach it.
163 confirmed := false
164 h := m.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
165 if !confirmed {
166 SetUser(r, "alice")
167 SetResource(r, "blog")
168 confirmed = true
169 }
170 }))
171
172 serve(h, http.MethodGet, "http://alice-blog.example.com/")
173 serve(h, http.MethodGet, "http://alice-blog.example.com/")
174 serve(h, http.MethodGet, "http://alice-ghost.example.com/")
175
176 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice", "resource": "blog"}); got != 2 {
177 t.Errorf("alice/blog requests = %v, want 2", got)
178 }
179 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice", "resource": ""}); got != 1 {
180 t.Errorf("alice requests without a resource = %v, want 1", got)
181 }
182 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "alice"}); got != 3 {
183 t.Errorf("alice requests = %v, want 3", got)
184 }
185}
186
187func TestUserMetricsMethod(t *testing.T) {
188 reg := prometheus.NewRegistry()
189 h := NewUserMetrics("test", reg, UserMetricsOpts{}).Handler(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
190
191 serve(h, "MADEUP", "/")
192
193 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"method": "OTHER"}); got != 1 {
194 t.Errorf("OTHER requests = %v, want 1", got)
195 }
196}
197
198func TestUserMetricsObserve(t *testing.T) {
199 reg := prometheus.NewRegistry()
200 m := NewUserMetrics("test", reg, UserMetricsOpts{})
201
202 m.Observe("carol", "app", http.StatusOK, http.MethodGet, 0, 7)
203 m.Observe("carol", "", http.StatusNotModified, "MADEUP", 0, -1)
204
205 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "carol", "status": "200", "method": "GET"}); got != 1 {
206 t.Errorf("200 requests = %v, want 1", got)
207 }
208 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "carol", "status": "304", "method": "OTHER"}); got != 1 {
209 t.Errorf("304 requests = %v, want 1", got)
210 }
211 if got := metricValue(t, reg, "pico_http_response_bytes_total", map[string]string{"user": "carol"}); got != 7 {
212 t.Errorf("bytes = %v, want 7", got)
213 }
214 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"user": "carol", "resource": "app"}); got != 1 {
215 t.Errorf("carol/app requests = %v, want 1", got)
216 }
217}
218
219func TestUserMetricsHijack(t *testing.T) {
220 reg := prometheus.NewRegistry()
221 h := NewUserMetrics("test", reg, UserMetricsOpts{}).Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
222 hj, ok := w.(http.Hijacker)
223 if !ok {
224 t.Error("writer does not implement http.Hijacker")
225 return
226 }
227 conn, rw, err := hj.Hijack()
228 if err != nil {
229 t.Errorf("hijack: %v", err)
230 return
231 }
232 _, _ = rw.WriteString("HTTP/1.1 101 Switching Protocols\r\n\r\n")
233 _ = rw.Flush()
234 _ = conn.Close()
235 }))
236 srv := httptest.NewServer(h)
237 defer srv.Close()
238
239 resp, err := http.Get(srv.URL)
240 if err != nil {
241 t.Fatalf("get: %v", err)
242 }
243 _ = resp.Body.Close()
244
245 if got := metricValue(t, reg, "pico_http_requests_total", map[string]string{"status": "101"}); got != 1 {
246 t.Errorf("101 requests = %v, want 1", got)
247 }
248}