3x-ui-agent/internal/web/service/online_store_test.go

package service

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

func sortedSet(in []string) map[string]bool {
	out := make(map[string]bool, len(in))
	for _, s := range in {
		out[s] = true
	}
	return out
}

func requireSet(t *testing.T, what string, got []string, want ...string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("%s: got %v, want set %v", what, got, want)
	}
	set := sortedSet(got)
	for _, w := range want {
		if !set[w] {
			t.Fatalf("%s: got %v, missing %q (want %v)", what, got, w, want)
		}
	}
}

// A node's reported tree must surface both in the flat union and in the
// GUID-keyed view, deduped against the local set and other nodes.
func TestOnlineStoreSetTreeUnionAndByGuid(t *testing.T) {
	st := &onlineStore{}
	st.SetNodeOnlineTree(1, map[string][]string{
		"guid-a": {"alice@example.com", "bob@example.com"},
		"guid-b": {"carol@example.com"},
	})
	st.SetNodeOnlineTree(2, map[string][]string{
		"guid-c": {"bob@example.com", "dave@example.com"},
	})
	// Local fold-in overlapping one node email (dedupe check).
	st.RefreshLocalOnline([]string{"alice@example.com", "local@example.com"}, nil, 1000, 20000)

	requireSet(t, "union", st.GetOnlineClients(),
		"alice@example.com", "bob@example.com", "carol@example.com", "dave@example.com", "local@example.com")

	merged := st.GetMergedNodeTrees()
	requireSet(t, "byGuid guid-a", merged["guid-a"], "alice@example.com", "bob@example.com")
	requireSet(t, "byGuid guid-b", merged["guid-b"], "carol@example.com")
	requireSet(t, "byGuid guid-c", merged["guid-c"], "bob@example.com", "dave@example.com")
	if _, ok := merged["guid-absent"]; ok {
		t.Fatalf("byGuid: unexpected key for absent guid: %v", merged)
	}
	requireSet(t, "local", st.GetLocalOnlineClients(), "alice@example.com", "local@example.com")
}

// Agent-node onlines arrive keyed by the node's attribution key (the
// synthNodeGuid form below for a GUID-less agent node — the controller used
// to file them under a literal "agent" key no inbound lookup ever matched);
// they must reach the union untouched.
func TestOnlineStoreAgentKeyTree(t *testing.T) {
	st := &onlineStore{}
	st.SetNodeOnlineTree(7, map[string][]string{"node:7": {"agentuser@example.com"}})
	requireSet(t, "union", st.GetOnlineClients(), "agentuser@example.com")
	requireSet(t, "byGuid node:7", st.GetMergedNodeTrees()["node:7"], "agentuser@example.com")
}

// Setting a tree replaces the node's previous contribution; an empty tree
// deletes it outright.
func TestOnlineStoreSetTreeReplacesAndEmptyDeletes(t *testing.T) {
	st := &onlineStore{}
	st.SetNodeOnlineTree(3, map[string][]string{"g": {"one@example.com"}})
	st.SetNodeOnlineTree(3, map[string][]string{"g": {"two@example.com"}})
	requireSet(t, "union after replace", st.GetOnlineClients(), "two@example.com")

	st.SetNodeOnlineTree(3, map[string][]string{})
	if got := st.GetOnlineClients(); len(got) != 0 {
		t.Fatalf("union after empty set: got %v, want empty", got)
	}
	if got := st.GetMergedNodeTrees(); len(got) != 0 {
		t.Fatalf("byGuid after empty set: got %v, want empty", got)
	}
}

// Active-inbound trees behave like online trees, with per-GUID dedupe.
func TestOnlineStoreActiveInboundTrees(t *testing.T) {
	st := &onlineStore{}
	st.SetNodeActiveInboundTree(1, map[string][]string{"g1": {"in-a", "in-b"}})
	st.SetNodeActiveInboundTree(2, map[string][]string{"g1": {"in-b", "in-c"}})
	merged := st.GetMergedActiveInboundTrees()
	requireSet(t, "active g1", merged["g1"], "in-a", "in-b", "in-c")

	st.RefreshLocalOnline(nil, []string{"local-in"}, 5000, 20000)
	requireSet(t, "local active", st.GetLocalActiveInbounds(), "local-in")
}

// Clear drops BOTH of the node's trees and leaves other nodes untouched.
func TestOnlineStoreClear(t *testing.T) {
	st := &onlineStore{}
	st.SetNodeOnlineTree(1, map[string][]string{"g1": {"a@example.com"}})
	st.SetNodeActiveInboundTree(1, map[string][]string{"g1": {"in-a"}})
	st.SetNodeOnlineTree(2, map[string][]string{"g2": {"b@example.com"}})

	st.ClearNodeOnlineClients(1)
	requireSet(t, "union after clear", st.GetOnlineClients(), "b@example.com")
	if got := st.GetMergedActiveInboundTrees(); len(got) != 0 {
		t.Fatalf("active after clear: got %v, want empty", got)
	}
}

// Local entries live for the grace window, then a prune-only refresh
// (nil, nil — what NodeTrafficSyncJob calls every tick) ages them out.
func TestOnlineStoreGraceWindowPrune(t *testing.T) {
	st := &onlineStore{}
	const grace = int64(20000)
	st.RefreshLocalOnline([]string{"stay@example.com"}, []string{"tag-stay"}, 10000, grace)
	st.RefreshLocalOnline([]string{"fresh@example.com"}, nil, 25000, grace)

	// At t=25000 the first entry is 15s old: still inside the window.
	requireSet(t, "local within grace", st.GetLocalOnlineClients(), "stay@example.com", "fresh@example.com")

	// Prune-only at t=45001: stay@10000 is 35s old (gone), fresh@25000 is
	// 20001ms old (also past the strict < grace comparison), tag pruned.
	st.RefreshLocalOnline(nil, nil, 45001, grace)
	if got := st.GetLocalOnlineClients(); len(got) != 0 {
		t.Fatalf("local after prune: got %v, want empty", got)
	}
	if got := st.GetLocalActiveInbounds(); len(got) != 0 {
		t.Fatalf("local active after prune: got %v, want empty", got)
	}
}

// Boundary: an entry exactly graceMs old is expired (original uses <).
func TestOnlineStoreGraceBoundary(t *testing.T) {
	st := &onlineStore{}
	st.RefreshLocalOnline([]string{"edge@example.com"}, nil, 0, 20000)
	st.RefreshLocalOnline(nil, nil, 19999, 20000)
	requireSet(t, "just inside", st.GetLocalOnlineClients(), "edge@example.com")
	st.RefreshLocalOnline(nil, nil, 20000, 20000)
	if got := st.GetLocalOnlineClients(); len(got) != 0 {
		t.Fatalf("at boundary: got %v, want empty", got)
	}
}

// The store is shared across InboundService values: a write through one
// instance must be visible through another (services are bare values; an
// instance field would fragment state between writers and readers).
func TestOnlineStoreSharedAcrossServiceInstances(t *testing.T) {
	writer := &InboundService{}
	reader := &InboundService{}
	const nodeID = 987654
	t.Cleanup(func() { writer.ClearNodeOnlineClients(nodeID) })

	writer.SetNodeOnlineTree(nodeID, map[string][]string{"guid-shared": {"shared@example.com"}})
	requireSet(t, "union via other instance", reader.GetOnlineClients(), "shared@example.com")
	byGuid := reader.GetOnlineClientsByGuid()
	requireSet(t, "byGuid via other instance", byGuid["guid-shared"], "shared@example.com")
}

// Concurrent writers/readers must not race or corrupt the union. Run with
// -race for the full guarantee; even without it, the final state is exact.
func TestOnlineStoreConcurrentAccess(t *testing.T) {
	st := &onlineStore{}
	const nodes = 8
	var wg sync.WaitGroup
	for n := 0; n < nodes; n++ {
		n := n
		wg.Add(1)
		go func() {
			defer wg.Done()
			guid := fmt.Sprintf("guid-%d", n)
			for i := 0; i < 50; i++ {
				st.SetNodeOnlineTree(n, map[string][]string{guid: {fmt.Sprintf("u%d-%d@example.com", n, i)}})
				_ = st.GetOnlineClients()
				_ = st.GetMergedNodeTrees()
				st.RefreshLocalOnline([]string{fmt.Sprintf("local-%d@example.com", n)}, nil, time.Now().UnixMilli(), 60000)
			}
		}()
	}
	wg.Wait()
	// Every node's last write holds exactly one email; locals may or may
	// not have been pruned by a same-ms refresh, so assert on node trees.
	merged := st.GetMergedNodeTrees()
	for n := 0; n < nodes; n++ {
		requireSet(t, fmt.Sprintf("tree %d", n), merged[fmt.Sprintf("guid-%d", n)],
			fmt.Sprintf("u%d-49@example.com", n))
	}
	for n := 0; n < nodes; n++ {
		st.ClearNodeOnlineClients(n)
	}
	if got := st.GetMergedNodeTrees(); len(got) != 0 {
		t.Fatalf("trees after mass clear: got %v, want empty", got)
	}
}