119 lines
4.4 KiB
Go
119 lines
4.4 KiB
Go
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)
|
|
}
|
|
}
|