3x-ui-agent/internal/web/job/node_heartbeat_agent_expiry_test.go

package job

import (
	"testing"
	"time"

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

func TestAgentHeartbeatStale(t *testing.T) {
	now := time.Unix(1_700_000_000, 0)
	ms := func(d time.Duration) int64 { return now.Add(-d).UnixMilli() }
	cases := []struct {
		name       string
		lastSeenMs int64
		want       bool
	}{
		{"never checked in", 0, false},
		{"negative stamp", -1, false},
		{"fresh heartbeat", ms(60 * time.Second), false},
		{"two missed heartbeats", ms(179 * time.Second), false},
		{"past threshold", ms(181 * time.Second), true},
		{"long gone", ms(24 * time.Hour), true},
	}
	for _, tc := range cases {
		if got := agentHeartbeatStale(tc.lastSeenMs, now); got != tc.want {
			t.Errorf("%s: agentHeartbeatStale(%d) = %v, want %v", tc.name, tc.lastSeenMs, got, tc.want)
		}
	}
}

func seedAgentNode(t *testing.T, name, status string, lastSeenMs int64) *model.Node {
	t.Helper()
	db := database.GetDB()
	n := &model.Node{
		Name:    name,
		Address: "agent",
		Port:    2053,
		Enable:  true,
		Status:  status,
	}
	if err := db.Create(n).Error; err != nil {
		t.Fatalf("create node: %v", err)
	}
	a := &model.AgentNode{
		TokenHash: "hash-" + name,
		Name:      name,
		NodeID:    &n.Id,
		LastSeen:  lastSeenMs,
	}
	if err := db.Create(a).Error; err != nil {
		t.Fatalf("create agent: %v", err)
	}
	return n
}

func nodeStatus(t *testing.T, id int) string {
	t.Helper()
	var n model.Node
	if err := database.GetDB().Where("id = ?", id).First(&n).Error; err != nil {
		t.Fatalf("reload node: %v", err)
	}
	return n.Status
}

func TestProbeOneExpiresStaleAgentNode(t *testing.T) {
	setupIntegrationDB(t)
	stale := time.Now().Add(-10 * time.Minute).UnixMilli()
	n := seedAgentNode(t, "stale-agent", "online", stale)

	(&NodeHeartbeatJob{}).probeOne(n)

	if got := nodeStatus(t, n.Id); got != "offline" {
		t.Fatalf("stale agent node status = %q, want offline", got)
	}
}

func TestProbeOneKeepsFreshAgentNodeOnline(t *testing.T) {
	setupIntegrationDB(t)
	fresh := time.Now().Add(-30 * time.Second).UnixMilli()
	n := seedAgentNode(t, "fresh-agent", "online", fresh)

	(&NodeHeartbeatJob{}).probeOne(n)

	if got := nodeStatus(t, n.Id); got != "online" {
		t.Fatalf("fresh agent node status = %q, want online", got)
	}
}

func TestProbeOneLeavesUnactivatedAgentNodeAlone(t *testing.T) {
	setupIntegrationDB(t)
	// Token issued, agent never checked in: last_seen = 0. The node must not
	// be forced offline — it was never online/activated in the first place.
	n := seedAgentNode(t, "dormant-agent", "unknown", 0)

	(&NodeHeartbeatJob{}).probeOne(n)

	if got := nodeStatus(t, n.Id); got != "unknown" {
		t.Fatalf("unactivated agent node status = %q, want unknown (untouched)", got)
	}
}

func TestProbeOnePanelNodeStillProbed(t *testing.T) {
	setupIntegrationDB(t)
	db := database.GetDB()
	// No agent bound: the legacy probe path must still run and mark an
	// unreachable panel node offline.
	n := &model.Node{
		Name:    "panel-node",
		Scheme:  "http",
		Address: "127.0.0.1",
		Port:    1, // nothing listens here; probe fails fast
		Enable:  true,
		Status:  "online",
	}
	if err := db.Create(n).Error; err != nil {
		t.Fatalf("create node: %v", err)
	}

	(&NodeHeartbeatJob{}).probeOne(n)

	if got := nodeStatus(t, n.Id); got != "offline" {
		t.Fatalf("unreachable panel node status = %q, want offline", got)
	}
}