feat(probe): port monitor-probe collector into agent, persist latest snapshot (item 7 B1)

2026-10-22 18:22:56 +00:00
parent 8e89960474
commit 0a973bc95c
6 changed files with 1496 additions and 3 deletions
@@ -0,0 +1,118 @@
+package controller
+
+import (
+	"bytes"
+	"net/http"
+	"net/http/httptest"
+	"path/filepath"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/database"
+	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
+)
+
+// Item 7 B1: the heartbeat/register optionally carry a monitor-probe-style
+// probe snapshot (facts + metrics). Reported fields land on the agent row;
+// a heartbeat without a probe block (old agents) must not clear stored
+// values, and a partial block only overwrites what it carries.
+func TestAgentHeartbeatProbe(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	dbDir := t.TempDir()
+	t.Setenv("XUI_DB_FOLDER", dbDir)
+	if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
+		t.Fatalf("InitDB: %v", err)
+	}
+	t.Cleanup(func() { _ = database.CloseDB() })
+	db := database.GetDB()
+
+	const token = "probe-test-token"
+	agent := &model.AgentNode{
+		TokenHash:   hashAgentToken(token),
+		TokenPrefix: "probe-",
+		Name:        "probe-node",
+	}
+	if err := db.Create(agent).Error; err != nil {
+		t.Fatalf("seed agent: %v", err)
+	}
+
+	engine := gin.New()
+	NewAgentController(engine.Group("/panel/api/agent"))
+
+	post := func(path, body string) int {
+		req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(body))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Authorization", "Bearer "+token)
+		rec := httptest.NewRecorder()
+		engine.ServeHTTP(rec, req)
+		return rec.Code
+	}
+	load := func() model.AgentNode {
+		var row model.AgentNode
+		if err := db.First(&row, agent.Id).Error; err != nil {
+			t.Fatalf("reload agent: %v", err)
+		}
+		return row
+	}
+
+	facts := `"probeFacts":{"hostname":"host-1","os":"Debian GNU/Linux 13","kernel":"6.12.1",` +
+		`"arch":"amd64","virt":"kvm","cpuName":"Test CPU","cpuCores":4,` +
+		`"memTotal":4000000000,"swapTotal":1000000000,"diskTotal":8000000000,` +
+		`"agentVersion":"A1.0.1","ipv4":"203.0.113.9","ipv6":"2001:db8::9"}`
+	metrics := `"probeMetrics":{"bootId":"boot-x/0123456789abcdef","iface":"",` +
+		`"uptime":3600,"cpu":12.5,"load":[0.5,0.25,0.125],` +
+		`"memTotal":4000000000,"memUsed":1500000000,"swapTotal":1000000000,"swapUsed":5000000,` +
+		`"diskTotal":8000000000,"diskUsed":3000000000,` +
+		`"netRxTotal":111000,"netTxTotal":222000,"netRx":1000,"netTx":2000,` +
+		`"tcp":42,"udp":7,"procs":123}`
+
+	// Register carries facts only.
+	if code := post("/panel/api/agent/register",
+		`{"token":"`+token+`","name":"probe-node",`+facts+`}`); code != http.StatusOK {
+		t.Fatalf("register with facts: status %d", code)
+	}
+	row := load()
+	if row.ProbeHostname != "host-1" || row.ProbeVirt != "kvm" || row.ProbeCPUCores != 4 ||
+		row.ProbeMemTotal != 4000000000 || row.ProbeDiskTotal != 8000000000 {
+		t.Fatalf("register facts not stored: %+v", row)
+	}
+	if row.ProbeAt != 0 {
+		t.Fatalf("register must not stamp ProbeAt (no metrics yet): %d", row.ProbeAt)
+	}
+
+	// Heartbeat with facts+metrics stores the latest snapshot.
+	if code := post("/panel/api/agent/heartbeat",
+		`{"token":"`+token+`","timestamp":1,`+facts+`,`+metrics+`}`); code != http.StatusOK {
+		t.Fatalf("heartbeat with probe: status %d", code)
+	}
+	row = load()
+	if row.ProbeBootID != "boot-x/0123456789abcdef" || row.ProbeMemUsed != 1500000000 ||
+		row.ProbeDiskUsed != 3000000000 || row.ProbeNetRxTot != 111000 || row.ProbeNetTxTot != 222000 ||
+		row.ProbeNetRx != 1000 || row.ProbeNetTx != 2000 || row.ProbeTCP != 42 || row.ProbeUDP != 7 ||
+		row.ProbeProcs != 123 || row.ProbeLoad1 != 0.5 || row.ProbeLoad15 != 0.125 || row.ProbeAt == 0 {
+		t.Fatalf("heartbeat probe not stored: %+v", row)
+	}
+	if row.ProbeCPU < 12.4 || row.ProbeCPU > 12.6 {
+		t.Fatalf("probe cpu = %v, want 12.5", row.ProbeCPU)
+	}
+
+	// Bare heartbeat (old agent): stored probe values survive untouched.
+	if code := post("/panel/api/agent/heartbeat", `{"token":"`+token+`","timestamp":2}`); code != http.StatusOK {
+		t.Fatalf("bare heartbeat: status %d", code)
+	}
+	row = load()
+	if row.ProbeBootID != "boot-x/0123456789abcdef" || row.ProbeMemUsed != 1500000000 || row.ProbeHostname != "host-1" {
+		t.Fatalf("bare heartbeat cleared probe: %+v", row)
+	}
+
+	// Partial metrics: only the carried field changes.
+	if code := post("/panel/api/agent/heartbeat",
+		`{"token":"`+token+`","timestamp":3,"probeMetrics":{"netRx":7777}}`); code != http.StatusOK {
+		t.Fatalf("partial heartbeat: status %d", code)
+	}
+	row = load()
+	if row.ProbeNetRx != 7777 || row.ProbeNetTx != 2000 || row.ProbeMemUsed != 1500000000 {
+		t.Fatalf("partial probe merged wrong: %+v", row)
+	}
+}