cloud-ip-validator/internal/probercore/probercore_test.go

package probercore

import (
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"sync"
	"sync/atomic"
	"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)
	}
}

// TestRegisterWithRetrySucceedsAfterSiteConfigured mirrors an admin adding
// this prober's site_id to control-api's config after the prober has
// already started: registration is rejected with 400 until then, and the
// prober must keep retrying rather than give up.
func TestRegisterWithRetrySucceedsAfterSiteConfigured(t *testing.T) {
	var attempts int32

	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/v1/probers/register", func(w http.ResponseWriter, r *http.Request) {
		n := atomic.AddInt32(&attempts, 1)
		if n < 3 {
			w.WriteHeader(http.StatusBadRequest)
			w.Write([]byte(`{"error":"unknown site_id: site-1"}`))
			return
		}
		w.WriteHeader(http.StatusOK)
	})
	ts := httptest.NewServer(mux)
	defer ts.Close()

	p := &Prober{
		cfg:                  &config.Prober{SiteID: "site-1", ControlAPIURL: ts.URL},
		client:               apiclient.New(ts.URL, 5*time.Second),
		log:                  testLogger(),
		registerRetryInitial: time.Millisecond,
		registerRetryMax:     5 * time.Millisecond,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	if err := p.registerWithRetry(ctx); err != nil {
		t.Fatalf("expected eventual success, got err: %v", err)
	}
	if got := atomic.LoadInt32(&attempts); got != 3 {
		t.Fatalf("expected exactly 3 attempts, got %d", got)
	}
}

// TestRegisterWithRetryStopsOnCancel confirms a prober waiting on an
// unreachable/unconfigured control-api can still be shut down promptly.
func TestRegisterWithRetryStopsOnCancel(t *testing.T) {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/v1/probers/register", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadRequest)
	})
	ts := httptest.NewServer(mux)
	defer ts.Close()

	p := &Prober{
		cfg:                  &config.Prober{SiteID: "site-1", ControlAPIURL: ts.URL},
		client:               apiclient.New(ts.URL, 5*time.Second),
		log:                  testLogger(),
		registerRetryInitial: 10 * time.Millisecond,
		registerRetryMax:     10 * time.Millisecond,
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(30 * time.Millisecond)
		cancel()
	}()

	var err error
	done := make(chan struct{})
	go func() {
		defer close(done)
		err = p.registerWithRetry(ctx)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("registerWithRetry did not return after ctx cancellation")
	}
	if err != context.Canceled {
		t.Fatalf("expected context.Canceled, got %v", err)
	}
}

// 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)
		}
	}
}

// TestWithTokenSendsBearerToControlAPI checks the prober's control-api
// client carries the agent token on every request.
func TestWithTokenSendsBearerToControlAPI(t *testing.T) {
	var mu sync.Mutex
	var auths []string
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		auths = append(auths, r.Header.Get("Authorization"))
		mu.Unlock()
		w.Write([]byte(`[]`))
	}))
	defer ts.Close()

	p := New(&config.Prober{SiteID: "site-1", ControlAPIURL: ts.URL}, testLogger()).WithToken("agent-secret")
	p.pollOnce(context.Background())

	mu.Lock()
	defer mu.Unlock()
	if len(auths) == 0 {
		t.Fatal("no requests reached control-api")
	}
	for _, a := range auths {
		if a != "Bearer agent-secret" {
			t.Fatalf("Authorization = %q, want bearer token", a)
		}
	}
}