130 lines
3.7 KiB
Go
130 lines
3.7 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)
|
|
}
|
|
}
|
|
|
|
// TestExtraChecksForPort confirms only ports 443 (TLS) and 22 (SSH) get an
|
|
// extra protocol check beyond the baseline TCPConnect every port gets.
|
|
func TestExtraChecksForPort(t *testing.T) {
|
|
// Use a host nothing is listening on so the extra check's own network
|
|
// call fails fast — only its CheckType is asserted here, the
|
|
// port->check mapping, not real connectivity.
|
|
closedHost := "127.0.0.1"
|
|
|
|
cases := []struct {
|
|
port int
|
|
wantLen int
|
|
wantCheckType string
|
|
}{
|
|
{port: 443, wantLen: 1, wantCheckType: "tls-443"},
|
|
{port: 22, wantLen: 1, wantCheckType: "ssh"},
|
|
{port: 80, wantLen: 0},
|
|
{port: 8080, wantLen: 0},
|
|
}
|
|
for _, tc := range cases {
|
|
extra := extraChecksForPort(closedHost, tc.port, 100*time.Millisecond)
|
|
if len(extra) != tc.wantLen {
|
|
t.Fatalf("port %d: expected %d extra checks, got %d", tc.port, tc.wantLen, len(extra))
|
|
}
|
|
if tc.wantLen == 0 {
|
|
continue
|
|
}
|
|
res := extra[0](context.Background())
|
|
if res.CheckType != tc.wantCheckType {
|
|
t.Fatalf("port %d: expected check type %q, got %q", tc.port, tc.wantCheckType, res.CheckType)
|
|
}
|
|
}
|
|
}
|