pangolin/server/internal/nodes/hub_test.go

package nodes_test

import (
	"context"
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"

	"github.com/wangjia/pangolin/server/internal/nodes"
	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

// newTestHub creates a Hub backed by a fresh miniredis instance.
func newTestHub(t *testing.T) (*nodes.Hub, *redis.Client, *miniredis.Miniredis) {
	t.Helper()
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	t.Cleanup(func() { rdb.Close() })
	return nodes.NewHub(rdb), rdb, mr
}

// ─── Push / Replay / Ack ─────────────────────────────────────────────────────

// TestHub_PushReplay verifies that pushed commands are retrievable via Replay,
// filtered by afterID (exclusive), and removed from the queue on Ack.
func TestHub_PushReplay(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()
	const nodeUUID = "node-replay-test"

	// Push three commands with explicit IDs so the test is deterministic.
	cmds := []*agentv1.Command{
		{CommandID: 10, Type: agentv1.CommandTypeUpsert,
			Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-a"}}},
		{CommandID: 20, Type: agentv1.CommandTypeUpsert,
			Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-b"}}},
		{CommandID: 30, Type: agentv1.CommandTypeRevoke,
			Revoke: &agentv1.RevokePayload{DpUUID: "dp-a"}},
	}
	for _, c := range cmds {
		if err := hub.Push(ctx, nodeUUID, c); err != nil {
			t.Fatalf("Push id=%d: %v", c.CommandID, err)
		}
	}

	// Replay from after 10 → must return 20 and 30 only.
	got, err := hub.Replay(ctx, nodeUUID, 10)
	if err != nil {
		t.Fatalf("Replay: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("Replay returned %d commands, want 2", len(got))
	}
	if got[0].CommandID != 20 || got[1].CommandID != 30 {
		t.Errorf("Replay ids = [%d %d], want [20 30]",
			got[0].CommandID, got[1].CommandID)
	}

	// Ack command 20.
	if err := hub.Ack(ctx, nodeUUID, 20); err != nil {
		t.Fatalf("Ack 20: %v", err)
	}

	// Replay from 0: 10 and 30 remain; 20 was removed.
	got, err = hub.Replay(ctx, nodeUUID, 0)
	if err != nil {
		t.Fatalf("Replay after ack: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("after Ack 20: Replay returned %d commands, want 2", len(got))
	}
	ids := []int64{got[0].CommandID, got[1].CommandID}
	if ids[0] != 10 || ids[1] != 30 {
		t.Errorf("after Ack 20: ids = %v, want [10 30]", ids)
	}
}

// TestHub_AutoID verifies that Hub assigns monotonically increasing IDs when
// cmd.CommandID == 0.
func TestHub_AutoID(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()
	const nodeUUID = "node-auto-id"

	c1 := &agentv1.Command{Type: agentv1.CommandTypeLifecycle,
		Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain}}
	c2 := &agentv1.Command{Type: agentv1.CommandTypeLifecycle,
		Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionResume}}

	if err := hub.Push(ctx, nodeUUID, c1); err != nil {
		t.Fatalf("Push c1: %v", err)
	}
	if err := hub.Push(ctx, nodeUUID, c2); err != nil {
		t.Fatalf("Push c2: %v", err)
	}

	if c1.CommandID == 0 {
		t.Error("c1.CommandID still 0 after Push")
	}
	if c2.CommandID == 0 {
		t.Error("c2.CommandID still 0 after Push")
	}
	if c1.CommandID >= c2.CommandID {
		t.Errorf("c1.CommandID=%d >= c2.CommandID=%d; want strictly increasing",
			c1.CommandID, c2.CommandID)
	}
}

// TestHub_AckIdempotent verifies that a second Ack for the same command is a no-op.
func TestHub_AckIdempotent(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()

	cmd := &agentv1.Command{CommandID: 99, Type: agentv1.CommandTypeUpsert,
		Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "x"}}}
	if err := hub.Push(ctx, "node-ack-idem", cmd); err != nil {
		t.Fatalf("Push: %v", err)
	}

	// First Ack: removes the command.
	if err := hub.Ack(ctx, "node-ack-idem", 99); err != nil {
		t.Fatalf("Ack 1: %v", err)
	}
	// Second Ack: must not return an error (ZREMRANGEBYSCORE with no match = 0 removed).
	if err := hub.Ack(ctx, "node-ack-idem", 99); err != nil {
		t.Errorf("Ack 2 (idempotent): unexpected error: %v", err)
	}

	// Queue should be empty.
	got, _ := hub.Replay(ctx, "node-ack-idem", 0)
	if len(got) != 0 {
		t.Errorf("after double Ack, queue has %d commands, want 0", len(got))
	}
}

// ─── Local delivery ───────────────────────────────────────────────────────────

// TestHub_LocalDelivery verifies that a Push reaches a registered subscriber's
// channel when the node is connected to the same hub instance.
func TestHub_LocalDelivery(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()
	const nodeUUID = "node-local-delivery"

	ch, done := hub.Register(nodeUUID)
	defer done()

	cmd := &agentv1.Command{
		Type:   agentv1.CommandTypeUpsert,
		Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-local"}},
	}
	if err := hub.Push(ctx, nodeUUID, cmd); err != nil {
		t.Fatalf("Push: %v", err)
	}

	select {
	case received := <-ch:
		if received.Upsert.Credential.DpUUID != "dp-local" {
			t.Errorf("got dp_uuid=%q, want dp-local", received.Upsert.Credential.DpUUID)
		}
	case <-time.After(time.Second):
		t.Fatal("timeout: command not delivered to local subscriber")
	}
}

// TestHub_RegisterUnregister verifies that after done() is called, the node no
// longer receives commands on the old channel.
func TestHub_RegisterUnregister(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()
	const nodeUUID = "node-unreg"

	ch, done := hub.Register(nodeUUID)
	done() // immediately unregister

	cmd := &agentv1.Command{
		CommandID: 42,
		Type:      agentv1.CommandTypeLifecycle,
		Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain},
	}
	if err := hub.Push(ctx, nodeUUID, cmd); err != nil {
		t.Fatalf("Push: %v", err)
	}

	select {
	case m := <-ch:
		// The channel might receive the command if it was buffered before unregister.
		// That is acceptable (it was still in the buffer at send time).
		// What must NOT happen is a panic or deadlock.
		_ = m
	case <-time.After(100 * time.Millisecond):
		// Expected: no delivery since we unregistered.
	}
}

// ─── Replay after reconnect ───────────────────────────────────────────────────

// TestHub_ReconnectResume verifies the "only resume unacked commands" invariant:
//   - Push commands 1, 2, 3
//   - Ack 1 and 2
//   - Reconnect with last_command_id = 3 (agent already has 3)
//   - Push command 4
//   - Replay from 3 → only command 4 delivered
func TestHub_ReconnectResume(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()
	const nodeUUID = "node-reconnect"

	push := func(id int64) {
		c := &agentv1.Command{
			CommandID: id,
			Type:      agentv1.CommandTypeLifecycle,
			Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain},
		}
		if err := hub.Push(ctx, nodeUUID, c); err != nil {
			t.Fatalf("Push %d: %v", id, err)
		}
	}

	push(1)
	push(2)
	push(3)

	// Ack 1 and 2 (simulating a previous session that processed them).
	for _, id := range []int64{1, 2} {
		if err := hub.Ack(ctx, nodeUUID, id); err != nil {
			t.Fatalf("Ack %d: %v", id, err)
		}
	}
	// Command 3 was delivered but not acked (stream dropped before ack).

	// Push command 4 (arrived while node was offline).
	push(4)

	// Reconnect: agent sends last_command_id = 3 (highest it processed).
	// Should receive command 3 (in queue, not acked) and command 4.
	got, err := hub.Replay(ctx, nodeUUID, 3)
	if err != nil {
		t.Fatalf("Replay: %v", err)
	}
	if len(got) != 1 {
		t.Fatalf("Replay from 3: got %d commands, want 1 (only cmd 4)", len(got))
	}
	if got[0].CommandID != 4 {
		t.Errorf("Replay from 3: got command_id=%d, want 4", got[0].CommandID)
	}

	// Reconnect with last_command_id = 2 (simulating agent that lost cmd 3):
	// should receive cmd 3 and 4.
	got, err = hub.Replay(ctx, nodeUUID, 2)
	if err != nil {
		t.Fatalf("Replay: %v", err)
	}
	if len(got) != 2 {
		t.Fatalf("Replay from 2: got %d commands, want 2 (cmds 3, 4)", len(got))
	}
	if got[0].CommandID != 3 || got[1].CommandID != 4 {
		t.Errorf("Replay from 2: ids = [%d, %d], want [3, 4]",
			got[0].CommandID, got[1].CommandID)
	}
}

// ─── Cross-instance pub/sub ───────────────────────────────────────────────────

// TestHub_CrossInstance starts two Hub instances connected to the same Redis,
// subscribes a node on hub1, pushes a command via hub2, and verifies delivery.
// This is the key validation that cross-instance command routing works.
func TestHub_CrossInstance(t *testing.T) {
	mr := miniredis.RunT(t)

	newHub := func() *nodes.Hub {
		rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
		t.Cleanup(func() { rdb.Close() })
		return nodes.NewHub(rdb)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	hub1 := newHub()
	hub2 := newHub()
	hub1.Start(ctx)
	hub2.Start(ctx)

	const nodeUUID = "node-cross-instance"
	ch, done := hub1.Register(nodeUUID)
	defer done()

	// Give the pub/sub goroutines a moment to subscribe to the Redis channel.
	time.Sleep(50 * time.Millisecond)

	cmd := &agentv1.Command{
		Type:   agentv1.CommandTypeUpsert,
		Upsert: &agentv1.UpsertPayload{Credential: &agentv1.Credential{DpUUID: "dp-cross"}},
	}
	if err := hub2.Push(ctx, nodeUUID, cmd); err != nil {
		t.Fatalf("hub2.Push: %v", err)
	}

	select {
	case received := <-ch:
		if received.Upsert.Credential.DpUUID != "dp-cross" {
			t.Errorf("got dp_uuid=%q, want dp-cross", received.Upsert.Credential.DpUUID)
		}
		t.Logf("cross-instance delivery: command_id=%d", received.CommandID)
	case <-time.After(3 * time.Second):
		t.Fatal("timeout: cross-instance command not delivered")
	}
}

// TestHub_CrossInstance_QueuePersisted verifies that a cross-instance push
// persists the command in Redis so Replay works even without a live subscriber.
func TestHub_CrossInstance_QueuePersisted(t *testing.T) {
	mr := miniredis.RunT(t)
	newHub := func() *nodes.Hub {
		rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
		t.Cleanup(func() { rdb.Close() })
		return nodes.NewHub(rdb)
	}

	ctx := context.Background()
	hub1 := newHub() // not started (no pubsub), no local subscriber
	hub2 := newHub()

	cmd := &agentv1.Command{
		Type:   agentv1.CommandTypeRevoke,
		Revoke: &agentv1.RevokePayload{DpUUID: "dp-persist"},
	}
	if err := hub2.Push(ctx, "offline-node", cmd); err != nil {
		t.Fatalf("hub2.Push: %v", err)
	}
	if cmd.CommandID == 0 {
		t.Fatal("command_id was not assigned")
	}

	// hub1 replays from the persistent queue (simulating reconnect).
	got, err := hub1.Replay(ctx, "offline-node", 0)
	if err != nil {
		t.Fatalf("hub1.Replay: %v", err)
	}
	if len(got) != 1 {
		t.Fatalf("got %d commands, want 1", len(got))
	}
	if got[0].Revoke.DpUUID != "dp-persist" {
		t.Errorf("got dp_uuid=%q, want dp-persist", got[0].Revoke.DpUUID)
	}
}

// TestHub_Broadcast delivers a command to multiple nodes.
func TestHub_Broadcast(t *testing.T) {
	hub, _, _ := newTestHub(t)
	ctx := context.Background()

	const n = 3
	nodes2 := make([]string, n)
	channels := make([]<-chan *agentv1.Command, n)
	dones := make([]func(), n)

	for i := 0; i < n; i++ {
		uuid := "node-bcast-" + string(rune('A'+i))
		nodes2[i] = uuid
		channels[i], dones[i] = hub.Register(uuid)
		defer dones[i]()
	}

	cmd := &agentv1.Command{
		Type:      agentv1.CommandTypeLifecycle,
		Lifecycle: &agentv1.LifecyclePayload{Action: agentv1.LifecycleActionDrain},
	}
	if err := hub.Broadcast(ctx, nodes2, cmd); err != nil {
		t.Fatalf("Broadcast: %v", err)
	}

	for i, ch := range channels {
		select {
		case received := <-ch:
			if received.Type != agentv1.CommandTypeLifecycle {
				t.Errorf("node %d: wrong type %v", i, received.Type)
			}
		case <-time.After(time.Second):
			t.Errorf("node %d: timeout waiting for broadcast", i)
		}
	}
}