1414package main
1515
1616import (
17+ "bytes"
1718 "errors"
1819 "fmt"
20+ "io"
1921 "net/http"
2022 "net/http/httptest"
23+ "net/url"
24+ "os"
25+ "strings"
26+ "sync"
2127 "testing"
28+ "time"
2229
30+ "github.com/cenkalti/backoff/v4"
2331 "github.com/prometheus/common/promslog"
2432)
2533
2634func prepareTest () (* httptest.Server , Coordinator ) {
35+ // This test server acts as the proxyURL
2736 ts := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
28- w .WriteHeader (http .StatusOK )
29- fmt .Fprintln (w , "GET /index.html HTTP/1.0\n \n OK" )
37+ switch r .URL .Path {
38+ case "/poll" :
39+ // On /poll, respond with an HTTP request serialized in the body
40+ var buf bytes.Buffer
41+ req , _ := http .NewRequest ("GET" , fmt .Sprintf ("http://%s/" , * myFqdn ), nil )
42+ req .Header .Set ("id" , "test-scrape-id" )
43+ req .Header .Set ("X-Prometheus-Scrape-Timeout-Seconds" , "10" )
44+ req .Write (& buf )
45+ w .WriteHeader (http .StatusOK )
46+ _ , _ = w .Write (buf .Bytes ())
47+ case "/push" :
48+ // Accept pushed scrape results, just respond OK
49+ io .Copy (io .Discard , r .Body )
50+ w .WriteHeader (http .StatusOK )
51+ default :
52+ w .WriteHeader (http .StatusNotFound )
53+ }
3054 }))
55+
3156 c := Coordinator {logger : promslog .NewNopLogger ()}
32- * proxyURL = ts .URL
57+ * proxyURL = ts .URL + "/"
58+ * myFqdn = "test.local" // Set fqdn to test.local for matching hostnames
59+
3360 return ts , c
3461}
3562
36- func TestDoScrape (t * testing.T ) {
63+ func TestDoScrape_Success (t * testing.T ) {
3764 ts , c := prepareTest ()
3865 defer ts .Close ()
3966
40- req , err := http .NewRequest ("GET" , ts .URL , nil )
67+ // Setup a test target server that will be scraped by doScrape
68+ targetServer := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
69+ // Verify Authorization header if set
70+ auth := r .Header .Get ("Authorization" )
71+ if auth != "" && auth != "Bearer dummy-token" {
72+ t .Errorf ("unexpected Authorization header: %s" , auth )
73+ }
74+ w .Header ().Set ("Content-Type" , "text/plain" )
75+ w .WriteHeader (http .StatusOK )
76+ fmt .Fprintln (w , "OK" )
77+ }))
78+ defer targetServer .Close ()
79+
80+ // Override myFqdn to match targetServer hostname
81+ u , err := url .Parse (targetServer .URL )
82+ if err != nil {
83+ t .Fatal (err )
84+ }
85+ * myFqdn = u .Hostname ()
86+
87+ // Prepare a scrape request targeting the test target server
88+ req , err := http .NewRequest ("GET" , targetServer .URL , nil )
89+ if err != nil {
90+ t .Fatal (err )
91+ }
92+ req .Header .Set ("id" , "scrape-id-123" )
93+ req .Header .Set ("X-Prometheus-Scrape-Timeout-Seconds" , "10" )
94+
95+ // Set bearerToken for authorization testing
96+ bearerTokenMutex .Lock ()
97+ bearerToken = "dummy-token"
98+ bearerTokenMutex .Unlock ()
99+
100+ c .doScrape (req , targetServer .Client ())
101+ }
102+
103+ func TestDoScrape_FailWrongFQDN (t * testing.T ) {
104+ ts , c := prepareTest ()
105+ defer ts .Close ()
106+
107+ req , err := http .NewRequest ("GET" , "http://wronghost.local" , nil )
41108 if err != nil {
42109 t .Fatal (err )
43110 }
44- req .Header .Add ("X-Prometheus-Scrape-Timeout-Seconds" , "10.0" )
45- * myFqdn = ts .URL
111+ req .Header .Set ("id" , "fail-id" )
112+ req .Header .Set ("X-Prometheus-Scrape-Timeout-Seconds" , "10" )
113+
114+ // This should cause handleErr due to fqdn mismatch
46115 c .doScrape (req , ts .Client ())
47116}
48117
@@ -57,10 +126,141 @@ func TestHandleErr(t *testing.T) {
57126 c .handleErr (req , ts .Client (), errors .New ("test error" ))
58127}
59128
60- func TestLoop (t * testing.T ) {
129+ func TestDoPush_ErrorOnInvalidProxyURL (t * testing.T ) {
130+ c := Coordinator {logger : promslog .NewNopLogger ()}
131+ * proxyURL = "http://%41:8080" // invalid URL (percent-encoding issue)
132+
133+ resp := & http.Response {
134+ StatusCode : http .StatusOK ,
135+ Body : io .NopCloser (strings .NewReader ("test" )),
136+ Header : http.Header {},
137+ }
138+ req , _ := http .NewRequest ("GET" , "http://example.com" , nil )
139+ err := c .doPush (resp , req , http .DefaultClient )
140+ if err == nil {
141+ t .Errorf ("expected error on invalid proxy URL, got nil" )
142+ }
143+ }
144+
145+ func TestDoPoll (t * testing.T ) {
61146 ts , c := prepareTest ()
62147 defer ts .Close ()
63- if err := c .doPoll (ts .Client ()); err != nil {
148+
149+ err := c .doPoll (ts .Client ())
150+ if err != nil {
151+ t .Fatalf ("doPoll failed: %v" , err )
152+ }
153+ }
154+
155+ func TestLoopWithBackoff (t * testing.T ) {
156+ var count int
157+ var mu sync.Mutex
158+ done := make (chan struct {})
159+ var once sync.Once
160+
161+ bo := backoffForTest (3 )
162+
163+ go func () {
164+ err := backoff .RetryNotify (func () error {
165+ mu .Lock ()
166+ defer mu .Unlock ()
167+ count ++
168+ if count > 2 {
169+ // safe close
170+ once .Do (func () { close (done ) })
171+ return errors .New ("forced error to stop retry" )
172+ }
173+ return errors .New ("temporary error" )
174+ }, bo , func (err error , d time.Duration ) {
175+ // No-op
176+ })
177+
178+ if err != nil {
179+ // safe even if already closed
180+ once .Do (func () { close (done ) })
181+ }
182+ }()
183+
184+ select {
185+ case <- done :
186+ case <- time .After (1 * time .Second ):
187+ t .Fatal ("loop test timed out" )
188+ }
189+ }
190+
191+
192+ func backoffForTest (maxRetries int ) backoff.BackOff {
193+ b := backoff .NewExponentialBackOff ()
194+ b .InitialInterval = 1 * time .Millisecond
195+ b .MaxInterval = 5 * time .Millisecond
196+ b .MaxElapsedTime = 10 * time .Millisecond
197+ return backoff .WithMaxRetries (b , uint64 (maxRetries ))
198+ }
199+
200+ func TestWatchBearerTokenFile (t * testing.T ) {
201+ // This function is hard to test fully without fsnotify events,
202+ // but we can test the initial loading of the token file.
203+
204+ // Create a temporary file with a token
205+ tmpfile := t .TempDir () + "/tokenfile"
206+ tokenContent := "file-token\n "
207+ if err := os .WriteFile (tmpfile , []byte (tokenContent ), 0600 ); err != nil {
208+ t .Fatal (err )
209+ }
210+
211+ logger := promslog .NewNopLogger ()
212+
213+ // Run watchBearerTokenFile in a goroutine; it will load token initially
214+ go func () {
215+ // This will block watching the directory, so we only wait shortly
216+ watchBearerTokenFile (tmpfile , logger )
217+ }()
218+
219+ // Wait briefly for the token to load
220+ time .Sleep (100 * time .Millisecond )
221+
222+ bearerTokenMutex .RLock ()
223+ defer bearerTokenMutex .RUnlock ()
224+ if bearerToken != strings .TrimSpace (tokenContent ) {
225+ t .Errorf ("expected bearer token %q, got %q" , strings .TrimSpace (tokenContent ), bearerToken )
226+ }
227+ }
228+
229+ func TestBearerTokenHeader (t * testing.T ) {
230+ token := "dummy-token"
231+ bearerTokenMutex .Lock ()
232+ bearerToken = token
233+ bearerTokenMutex .Unlock ()
234+
235+ var receivedToken string
236+
237+ ts := httptest .NewServer (http .HandlerFunc (func (w http.ResponseWriter , r * http.Request ) {
238+ receivedToken = r .Header .Get ("Authorization" )
239+ w .WriteHeader (http .StatusOK )
240+ }))
241+ defer ts .Close ()
242+
243+ // Ensure myFqdn matches the test server's hostname
244+ u , err := url .Parse (ts .URL )
245+ if err != nil {
246+ t .Fatal (err )
247+ }
248+ * myFqdn = u .Hostname ()
249+
250+ req , err := http .NewRequest ("GET" , ts .URL , nil )
251+ if err != nil {
64252 t .Fatal (err )
65253 }
254+
255+ // Set required headers for doScrape to accept this request
256+ req .Header .Set ("id" , "token-test-id" )
257+ req .Header .Set ("X-Prometheus-Scrape-Timeout-Seconds" , "10" )
258+
259+ c := Coordinator {logger : promslog .NewNopLogger ()}
260+ c .doScrape (req , ts .Client ())
261+
262+ expected := "Bearer dummy-token"
263+ if receivedToken != expected {
264+ t .Fatalf ("expected %q, got %q" , expected , receivedToken )
265+ }
66266}
0 commit comments