3x-ui-agent/internal/agent/fast_report_test.go

package agent

import (
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
)

// V2.0.1 probe batch, agent side: the config poll's cadence + pin
// are adopted; while the fast channel is active the heartbeat stops
// carrying metrics (they ride the fast report instead); a master
// without the endpoint (404) sends the agent back to heartbeat
// metrics permanently.

type fastMaster struct {
	mu           sync.Mutex
	heartbeat    []byte
	fastReport   []byte
	fastStatus   int
	configBody   string
	reportStatus int
}

func (m *fastMaster) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		m.mu.Lock()
		switch {
		case strings.HasSuffix(r.URL.Path, "/config"):
			m.mu.Unlock()
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(m.configBody))
		case strings.HasSuffix(r.URL.Path, "/probe-report"):
			m.fastReport = body
			st := m.fastStatus
			m.mu.Unlock()
			if st != 0 && st != http.StatusOK {
				w.WriteHeader(st)
				return
			}
			_, _ = w.Write([]byte(`{"success":true}`))
		case strings.HasSuffix(r.URL.Path, "/heartbeat"):
			m.heartbeat = body
			m.mu.Unlock()
			_, _ = w.Write([]byte(`{"success":true}`))
		default:
			m.mu.Unlock()
			http.NotFound(w, r)
		}
	})
}

func TestFastReportSplitAndFallback(t *testing.T) {
	SetStatIfacePin("")
	t.Cleanup(func() { SetStatIfacePin("") })

	master := &fastMaster{configBody: `{"probe_fast_interval":5,"probe_stat_iface":"eth7"}`}
	srv := httptest.NewServer(master.handler())
	defer srv.Close()

	a := &Agent{
		config:     &Config{MasterURL: srv.URL, Token: "tok", NodeName: "n"},
		httpClient: http.DefaultClient,
		probes:     newProbeRunner(),
	}

	if err := a.fetchAndApplyConfig(); err != nil {
		t.Fatalf("fetchAndApplyConfig: %v", err)
	}
	if got := a.fastIntervalSec.Load(); got != 5 {
		t.Fatalf("fast interval = %d, want 5 from the poll", got)
	}
	if got := statIfacePinValue(); got != "eth7" {
		t.Fatalf("stat iface pin = %q, want eth7", got)
	}
	if !a.fastReportActive() {
		t.Fatal("fast report should be active after the poll offered a cadence")
	}

	// The fast report carries the metrics…
	a.sendFastReport()
	master.mu.Lock()
	fast := master.fastReport
	master.mu.Unlock()
	var fastBody map[string]json.RawMessage
	if err := json.Unmarshal(fast, &fastBody); err != nil || fastBody["probeMetrics"] == nil {
		t.Fatalf("fast report = %s (err %v), want a probeMetrics block", fast, err)
	}

	// …and the heartbeat no longer does.
	if err := a.sendHeartbeat(); err != nil {
		t.Fatalf("sendHeartbeat: %v", err)
	}
	master.mu.Lock()
	hb := master.heartbeat
	master.mu.Unlock()
	var hbBody map[string]json.RawMessage
	if err := json.Unmarshal(hb, &hbBody); err != nil {
		t.Fatalf("heartbeat body: %v", err)
	}
	if hbBody["probeMetrics"] != nil {
		t.Error("heartbeat still carries probeMetrics while the fast channel is active")
	}
	if hbBody["probeFacts"] == nil || hbBody["telemetry"] == nil {
		t.Error("heartbeat lost its heavy content (facts/telemetry)")
	}

	// A 404 from the fast endpoint (old master) disables the fast
	// channel and restores heartbeat metrics.
	master.mu.Lock()
	master.fastStatus = http.StatusNotFound
	master.mu.Unlock()
	a.sendFastReport()
	if a.fastReportActive() {
		t.Fatal("fast report still active after a 404")
	}
	if err := a.sendHeartbeat(); err != nil {
		t.Fatalf("sendHeartbeat after fallback: %v", err)
	}
	master.mu.Lock()
	hb = master.heartbeat
	master.mu.Unlock()
	hbBody = map[string]json.RawMessage{}
	if err := json.Unmarshal(hb, &hbBody); err != nil {
		t.Fatalf("heartbeat body after fallback: %v", err)
	}
	if hbBody["probeMetrics"] == nil {
		t.Error("heartbeat did not resume carrying probeMetrics after the 404 fallback")
	}
}

// The counted-interface pin: strict single-interface selection, and
// a pin naming a vanished interface counts nothing (so the epoch
// changes and the master re-baselines instead of booking garbage).
func TestPinnedInterfaces(t *testing.T) {
	all := []netCounters{
		{name: "eth0", rx: 100, tx: 200},
		{name: "ens3", rx: 300, tx: 400},
	}
	got := pinnedInterfaces(all, "ens3")
	if len(got) != 1 || got[0].name != "ens3" || got[0].rx != 300 {
		t.Errorf("pinnedInterfaces = %+v, want only ens3", got)
	}
	if got := pinnedInterfaces(all, "gone0"); len(got) != 0 {
		t.Errorf("pin for a vanished interface = %+v, want empty", got)
	}
}