3x-ui-agent/internal/web/controller/probe_tasks_agent_test.go

package controller

import (
	"bytes"
	"encoding/json"
	"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"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"
)

// Item 7 B4: the agent-facing wire — probe_tasks rides the config
// response (always present, empty list clears), and probeResults on
// the heartbeat are filed stamped at arrival, with junk rows
// (unknown/absent latency, non-positive task ids) skipped and
// unassigned tasks dropped by the store's assignment gate.
func TestAgentProbeTasksWire(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()

	node := &model.Node{Name: "b4-wire"}
	if err := db.Create(node).Error; err != nil {
		t.Fatal(err)
	}
	const token = "b4-wire-token"
	agent := &model.AgentNode{
		TokenHash: hashAgentToken(token), TokenPrefix: "b4-",
		Name: "b4-wire-agent", NodeID: &node.Id,
	}
	if err := db.Create(agent).Error; err != nil {
		t.Fatal(err)
	}
	agentID := int64(agent.Id)
	probe.DeleteAgentData(agentID)

	taskID, err := probe.SavePingTask(probe.PingTask{
		Name: "wire", Target: "1.1.1.1:443", Interval: 30,
		AgentIDs: []int64{agentID},
	}, nil)
	if err != nil {
		t.Fatal(err)
	}

	engine := gin.New()
	NewAgentController(engine.Group("/panel/api/agent"))
	do := func(method, path, body string) (int, []byte) {
		var rd *bytes.Reader
		if body != "" {
			rd = bytes.NewReader([]byte(body))
		} else {
			rd = bytes.NewReader(nil)
		}
		req := httptest.NewRequest(method, path, rd)
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+token)
		rec := httptest.NewRecorder()
		engine.ServeHTTP(rec, req)
		return rec.Code, rec.Body.Bytes()
	}

	const cfgPath = "/panel/api/agent/config"
	const hbPath = "/panel/api/agent/heartbeat"

	// Config carries the task set.
	code, body := do(http.MethodGet, cfgPath, "")
	if code != http.StatusOK {
		t.Fatalf("config: %d %s", code, body)
	}
	var cfg struct {
		ProbeTasks []probe.PingTaskRef `json:"probe_tasks"`
	}
	if err := json.Unmarshal(body, &cfg); err != nil {
		t.Fatal(err)
	}
	if len(cfg.ProbeTasks) != 1 || cfg.ProbeTasks[0].ID != taskID ||
		cfg.ProbeTasks[0].Target != "1.1.1.1:443" || cfg.ProbeTasks[0].Interval != 30 {
		t.Fatalf("config probe_tasks: %s", body)
	}

	// Heartbeat with a mixed batch: only the valid, assigned sample
	// lands. (All rows share the arrival stamp, so the two filings of
	// the same task collapse — assert on the surviving value.)
	hb := `{"probeResults":[` +
		`{"taskId":` + itoa(taskID) + `,"latencyMs":42},` +
		`{"taskId":999999,"latencyMs":5},` +
		`{"taskId":-3,"latencyMs":5},` +
		`{"taskId":` + itoa(taskID) + `}` +
		`]}`
	code, body = do(http.MethodPost, hbPath, hb)
	if code != http.StatusOK {
		t.Fatalf("heartbeat: %d %s", code, body)
	}
	st := probe.Default()
	var count int64
	var latency int64
	if err := st.DB().QueryRow(`SELECT COUNT(*), COALESCE(MAX(latency), -999) FROM probe_ping_record WHERE agent_id = ?`, agentID).Scan(&count, &latency); err != nil {
		t.Fatal(err)
	}
	if count != 1 || latency != 42 {
		t.Fatalf("filed rows %d latency %d", count, latency)
	}

	// After the task is unassigned, config reports an empty (present)
	// list and further results are dropped by the assignment gate.
	if err := probe.SetNodePingTasks(agentID, []int64{}, []int64{taskID}); err != nil {
		t.Fatal(err)
	}
	code, body = do(http.MethodGet, cfgPath, "")
	if code != http.StatusOK {
		t.Fatalf("config 2: %d", code)
	}
	if !bytes.Contains(body, []byte(`"probe_tasks":[]`)) {
		t.Fatalf("empty probe_tasks not present: %s", body)
	}
	code, _ = do(http.MethodPost, hbPath, `{"probeResults":[{"taskId":`+itoa(taskID)+`,"latencyMs":77}]}`)
	if code != http.StatusOK {
		t.Fatalf("heartbeat 2: %d", code)
	}
	if err := st.DB().QueryRow(`SELECT COUNT(*) FROM probe_ping_record WHERE agent_id = ? AND latency = 77`, agentID).Scan(&count); err != nil {
		t.Fatal(err)
	}
	if count != 0 {
		t.Fatal("unassigned result was filed")
	}
	probe.DeletePingTask(taskID)
}