97 lines
2.6 KiB
Go
97 lines
2.6 KiB
Go
package probercore
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"cloudipvalidator/internal/apiclient"
|
|
"cloudipvalidator/internal/config"
|
|
)
|
|
|
|
func testLogger() *slog.Logger {
|
|
return slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
|
|
}
|
|
|
|
// TestPollOnceCallsHeartbeatBeforeAssignments mirrors agentcore's
|
|
// heartbeat-then-assignment ordering: the prober must announce it's alive
|
|
// on every poll cycle before fetching work, not just once at startup.
|
|
func TestPollOnceCallsHeartbeatBeforeAssignments(t *testing.T) {
|
|
var mu sync.Mutex
|
|
var calls []string
|
|
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("POST /api/v1/probers/site-1/heartbeat", func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
calls = append(calls, "heartbeat")
|
|
mu.Unlock()
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write([]byte(`{"ok":true}`))
|
|
})
|
|
mux.HandleFunc("GET /api/v1/probers/site-1/assignments", func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
calls = append(calls, "assignments")
|
|
mu.Unlock()
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write([]byte(`[]`))
|
|
})
|
|
ts := httptest.NewServer(mux)
|
|
defer ts.Close()
|
|
|
|
p := &Prober{
|
|
cfg: &config.Prober{SiteID: "site-1"},
|
|
client: apiclient.New(ts.URL, 5*time.Second),
|
|
log: testLogger(),
|
|
}
|
|
p.pollOnce(context.Background())
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(calls) != 2 || calls[0] != "heartbeat" || calls[1] != "assignments" {
|
|
t.Fatalf("expected [heartbeat assignments], got %v", calls)
|
|
}
|
|
}
|
|
|
|
// TestPollOnceSkipsAssignmentsWhenHeartbeatFails confirms a failed
|
|
// heartbeat aborts the poll cycle early — mirrors agentcore.pollOnce's
|
|
// same early-return behavior.
|
|
func TestPollOnceSkipsAssignmentsWhenHeartbeatFails(t *testing.T) {
|
|
var mu sync.Mutex
|
|
var calls []string
|
|
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("POST /api/v1/probers/site-1/heartbeat", func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
calls = append(calls, "heartbeat")
|
|
mu.Unlock()
|
|
w.WriteHeader(http.StatusNotFound)
|
|
})
|
|
mux.HandleFunc("GET /api/v1/probers/site-1/assignments", func(w http.ResponseWriter, r *http.Request) {
|
|
mu.Lock()
|
|
calls = append(calls, "assignments")
|
|
mu.Unlock()
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write([]byte(`[]`))
|
|
})
|
|
ts := httptest.NewServer(mux)
|
|
defer ts.Close()
|
|
|
|
p := &Prober{
|
|
cfg: &config.Prober{SiteID: "site-1"},
|
|
client: apiclient.New(ts.URL, 5*time.Second),
|
|
log: testLogger(),
|
|
}
|
|
p.pollOnce(context.Background())
|
|
|
|
mu.Lock()
|
|
defer mu.Unlock()
|
|
if len(calls) != 1 || calls[0] != "heartbeat" {
|
|
t.Fatalf("expected only [heartbeat] (assignments skipped on heartbeat failure), got %v", calls)
|
|
}
|
|
}
|