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

package node_test

import (
	"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 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:    "oto-001",
		Alias:     "oto",
		AgentKind: config.AgentKindOTOAgent,
	})

	entry, ok := reg.Get("oto-001")
	if !ok {
		t.Fatal("expected oto entry to be registered")
	}
	if entry.AgentKind != config.AgentKindOTOAgent {
		t.Fatalf("agent kind: got %q want %q", entry.AgentKind, config.AgentKindOTOAgent)
	}
	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)
	}
}