iop/apps/edge/internal/node/registry_test.go

package node_test

import (
	"sync"
	"sync/atomic"
	"testing"

	toki "git.toki-labs.com/toki/proto-socket/go"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

func TestRegistry_RegisterAndCount(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID: "node-001",
		Alias:  "local-node",
		Client: &toki.TcpClient{},
	}
	reg.Register(entry)
	if reg.Count() != 1 {
		t.Fatalf("expected 1 node, got %d", reg.Count())
	}
}

func TestRegistryRegisterIfAbsentRejectsConcurrentDuplicate(t *testing.T) {
	reg := edgenode.NewRegistry()

	var accepted int32
	var wg sync.WaitGroup
	start := make(chan struct{})
	for i := 0; i < 32; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			if reg.RegisterIfAbsent(&edgenode.NodeEntry{
				NodeID: "node-dup",
				Alias:  "dup",
				Client: &toki.TcpClient{},
			}) {
				atomic.AddInt32(&accepted, 1)
			}
		}()
	}
	close(start)
	wg.Wait()

	if accepted != 1 {
		t.Fatalf("accepted registrations: got %d want 1", accepted)
	}
	if reg.Count() != 1 {
		t.Fatalf("registry count: got %d want 1", reg.Count())
	}
}

func TestRegistryUnregisterIfClientIgnoresStaleConnection(t *testing.T) {
	reg := edgenode.NewRegistry()
	live := &toki.TcpClient{}
	stale := &toki.TcpClient{}
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1", Client: live})

	if reg.UnregisterIfClient("node-1", stale) {
		t.Fatal("stale client should not unregister live entry")
	}
	entry, ok := reg.Get("node-1")
	if !ok {
		t.Fatal("live entry was removed by stale client")
	}
	if entry.Client != live {
		t.Fatal("live entry client changed unexpectedly")
	}

	if !reg.UnregisterIfClient("node-1", live) {
		t.Fatal("live client should unregister entry")
	}
	if reg.Count() != 0 {
		t.Fatalf("registry count after live unregister: got %d want 0", reg.Count())
	}
}

func TestRegistry_Resolve_ByAliasOrID(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{
		NodeID: "node-001",
		Alias:  "local-node",
	}
	reg.Register(entry)

	// Resolve by ID
	if e, err := reg.Resolve("node-001"); err != nil || e.NodeID != "node-001" {
		t.Errorf("failed to resolve by ID: %v", err)
	}

	// Resolve by Alias
	if e, err := reg.Resolve("local-node"); err != nil || e.NodeID != "node-001" {
		t.Errorf("failed to resolve by Alias: %v", err)
	}

	// Resolve non-existent
	if _, err := reg.Resolve("unknown"); err == nil {
		t.Error("expected error for unknown node")
	}
}

func TestRegistry_Resolve_ByDisplayLabel(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-001", Alias: "local-node"})
	reg.Register(&edgenode.NodeEntry{NodeID: "node-002", Alias: "remote-node"})

	if e, err := reg.Resolve("node0"); err != nil || e.NodeID != "node-001" {
		t.Fatalf("Resolve(node0) = %+v, %v; want node-001", e, err)
	}
	if e, err := reg.Resolve("node1"); err != nil || e.NodeID != "node-002" {
		t.Fatalf("Resolve(node1) = %+v, %v; want node-002", e, err)
	}
}

func TestRegistry_Resolve_ImplicitSingleNodeOnly(t *testing.T) {
	reg := edgenode.NewRegistry()

	// Empty
	if _, err := reg.Resolve(""); err == nil {
		t.Error("expected error on empty registry")
	}

	// Single node
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
	if e, err := reg.Resolve(""); err != nil || e.NodeID != "node-1" {
		t.Errorf("failed implicit resolve for single node: %v", err)
	}

	// Multiple nodes
	reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
	if _, err := reg.Resolve(""); err == nil {
		t.Error("expected error for implicit resolve with multiple nodes")
	}
}

func TestRegistryRegisterPreservesAgentKind(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:    "node-explicit",
		Alias:     "explicit",
		AgentKind: config.AgentKindGenericNode,
	})

	entry, ok := reg.Get("node-explicit")
	if !ok {
		t.Fatal("expected explicit entry to be registered")
	}
	if entry.AgentKind != config.AgentKindGenericNode {
		t.Fatalf("agent kind: got %q want %q", entry.AgentKind, config.AgentKindGenericNode)
	}
	if entry.LifecycleState != edgenode.LifecycleConnected {
		t.Fatalf("lifecycle: got %q want %q", entry.LifecycleState, edgenode.LifecycleConnected)
	}
}

func TestRegistryRegisterDefaultsKindAndLifecycle(t *testing.T) {
	reg := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{NodeID: "node-001", Alias: "local-node"}
	reg.Register(entry)

	if entry.AgentKind != config.AgentKindGenericNode {
		t.Fatalf("agent kind default: got %q want %q", entry.AgentKind, config.AgentKindGenericNode)
	}
	if entry.LifecycleState != edgenode.LifecycleConnected {
		t.Fatalf("lifecycle default: got %q want %q", entry.LifecycleState, edgenode.LifecycleConnected)
	}
}

func TestRegistry_UnregisterRemovesAlias(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})

	reg.Unregister("node-1")
	if reg.Count() != 0 {
		t.Fatalf("expected 0 nodes, got %d", reg.Count())
	}

	if _, err := reg.Resolve("alias-1"); err == nil {
		t.Error("alias should be removed after unregister")
	}
}

func TestRegistryUpdateLifecycle(t *testing.T) {
	reg := edgenode.NewRegistry()

	if ok := reg.UpdateLifecycle("non-existent", edgenode.LifecycleOnline); ok {
		t.Error("expected false for updating non-existent node")
	}

	entry := &edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"}
	reg.Register(entry)

	if ok := reg.UpdateLifecycle("node-1", edgenode.LifecycleOnline); !ok {
		t.Error("expected true for updating existing node")
	}

	retrieved, ok := reg.Get("node-1")
	if !ok {
		t.Fatal("failed to get node")
	}
	if retrieved.LifecycleState != edgenode.LifecycleOnline {
		t.Errorf("lifecycle: got %q, want %q", retrieved.LifecycleState, edgenode.LifecycleOnline)
	}
}