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

package controller

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
)

// Item 20 (master side): a pending update intent rides the agent's config
// poll; outcome reports come back on heartbeats; a heartbeat whose
// telemetry reaches the target version closes the intent as successful.

const updToken = "update-status-token-0000000000000000000001"

func setupUpdateStatusTest(t *testing.T) (*gin.Engine, *model.AgentNode) {
	t.Helper()
	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: "upd-node", Address: "10.0.9.3"}
	if err := db.Create(node).Error; err != nil {
		t.Fatalf("seed node: %v", err)
	}
	agent := &model.AgentNode{
		TokenHash:         hashAgentToken(updToken),
		TokenPrefix:       updToken[:12],
		Name:              "upd-node",
		NodeID:            &node.Id,
		AgentVersion:      "A1.0.0",
		UpdateRequestedAt: 1700000000000,
		UpdateTarget:      "A1.0.1",
		UpdateTag:         "V1.0.1",
		UpdateState:       "pending",
	}
	if err := db.Create(agent).Error; err != nil {
		t.Fatalf("seed agent: %v", err)
	}
	engine := gin.New()
	NewAgentController(engine.Group("/panel/api/agent"))
	return engine, agent
}

func postHeartbeat(t *testing.T, engine *gin.Engine, body string) {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/panel/api/agent/heartbeat", bytes.NewBufferString(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+updToken)
	rec := httptest.NewRecorder()
	engine.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("heartbeat: status %d", rec.Code)
	}
}

func TestConfigCarriesPendingUpdateIntent(t *testing.T) {
	engine, _ := setupUpdateStatusTest(t)
	req := httptest.NewRequest(http.MethodGet, "/panel/api/agent/config", nil)
	req.Header.Set("Authorization", "Bearer "+updToken)
	rec := httptest.NewRecorder()
	engine.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("config: status %d body %s", rec.Code, rec.Body.String())
	}
	var resp struct {
		AgentUpdate *struct {
			RequestedAt int64  `json:"requested_at"`
			Target      string `json:"target"`
			Tag         string `json:"tag"`
			BaseURL     string `json:"base_url"`
		} `json:"agent_update"`
	}
	if err := json.Unmarshal(rec.Body.Bytes(), &resp); err != nil {
		t.Fatalf("config response not JSON: %v", err)
	}
	if resp.AgentUpdate == nil {
		t.Fatal("pending intent missing from config response")
	}
	if resp.AgentUpdate.Target != "A1.0.1" || resp.AgentUpdate.Tag != "V1.0.1" ||
		resp.AgentUpdate.RequestedAt != 1700000000000 ||
		!strings.HasSuffix(resp.AgentUpdate.BaseURL, "/") {
		t.Fatalf("intent payload wrong: %+v", resp.AgentUpdate)
	}
}

func TestHeartbeatFailedReportClearsIntent(t *testing.T) {
	engine, agent := setupUpdateStatusTest(t)
	postHeartbeat(t, engine, `{"updateStatus":{"requestedAt":1700000000000,"state":"failed","error":"sha256 mismatch"}}`)

	row := reloadAgent(t, agent.Id)
	if row.UpdateState != "failed" || row.UpdateError != "sha256 mismatch" || row.UpdateRequestedAt != 0 {
		t.Fatalf("failed report not folded: %+v", row)
	}

	// A stale report for the cleared intent must not resurrect it.
	postHeartbeat(t, engine, `{"updateStatus":{"requestedAt":1700000000000,"state":"downloading"}}`)
	row = reloadAgent(t, agent.Id)
	if row.UpdateState != "failed" {
		t.Fatalf("stale report changed state: %+v", row)
	}
}

func TestHeartbeatNewVersionConfirmsSuccess(t *testing.T) {
	engine, agent := setupUpdateStatusTest(t)
	postHeartbeat(t, engine, `{"telemetry":{"agentVersion":"A1.0.1"}}`)

	row := reloadAgent(t, agent.Id)
	if row.UpdateState != "success" || row.UpdateRequestedAt != 0 {
		t.Fatalf("success not confirmed: %+v", row)
	}
}