yao/tai/registry/registry_test.go

package registry

import (
	"log/slog"
	"sync"
	"testing"
	"time"

	"github.com/yaoapp/yao/tai/types"
)

// newTestRegistry creates a standalone registry for testing (bypasses global singleton).
func newTestRegistry() *Registry {
	return &Registry{
		nodes:  make(map[string]*TaiNode),
		logger: slog.Default(),
	}
}

func TestRegister_SetsFieldsAndOnline(t *testing.T) {
	r := newTestRegistry()
	node := &TaiNode{
		TaiID:     "tai-001",
		MachineID: "m-abc",
		Version:   "1.0.0",
		Mode:      "tunnel",
		Ports:     types.Ports{GRPC: 19100},
	}
	r.Register(node)

	snap, ok := r.Get("tai-001")
	if !ok {
		t.Fatal("expected node to exist after Register")
	}
	if snap.Status != "online" {
		t.Errorf("Status = %q, want online", snap.Status)
	}
	if snap.MachineID != "m-abc" {
		t.Errorf("MachineID = %q, want m-abc", snap.MachineID)
	}
	if snap.ConnectedAt.IsZero() {
		t.Error("ConnectedAt should be set")
	}
}

func TestRegister_Overwrite(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-001", Version: "1.0"})
	r.Register(&TaiNode{TaiID: "tai-001", Version: "2.0"})

	snap, ok := r.Get("tai-001")
	if !ok {
		t.Fatal("node should exist")
	}
	if snap.Version != "2.0" {
		t.Errorf("Version = %q, want 2.0 after re-register", snap.Version)
	}
}

func TestUnregister_RemovesNode(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-001"})
	r.Unregister("tai-001")

	if _, ok := r.Get("tai-001"); ok {
		t.Error("expected node to be removed after Unregister")
	}
}

func TestUnregister_Nonexistent(t *testing.T) {
	r := newTestRegistry()
	r.Unregister("ghost")
}

func TestGet_NotFound(t *testing.T) {
	r := newTestRegistry()
	if _, ok := r.Get("missing"); ok {
		t.Error("expected false for missing node")
	}
}

func TestList_Empty(t *testing.T) {
	r := newTestRegistry()
	if got := r.List(); len(got) != 0 {
		t.Errorf("List() = %d items, want 0", len(got))
	}
}

func TestList_MultipleNodes(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "a"})
	r.Register(&TaiNode{TaiID: "b"})
	r.Register(&TaiNode{TaiID: "c"})

	list := r.List()
	if len(list) != 3 {
		t.Errorf("List() = %d items, want 3", len(list))
	}

	ids := map[string]bool{}
	for _, snap := range list {
		ids[snap.TaiID] = true
	}
	for _, id := range []string{"a", "b", "c"} {
		if !ids[id] {
			t.Errorf("missing node %q in List()", id)
		}
	}
}

func TestSnapshot_DeepCopy(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{
		TaiID: "tai-001",
		Ports: types.Ports{GRPC: 19100, HTTP: 8099},
	})

	snap, _ := r.Get("tai-001")
	snap.Ports.GRPC = 0

	snap2, _ := r.Get("tai-001")
	if snap2.Ports.GRPC != 19100 {
		t.Error("snapshot modification leaked into registry node")
	}
}

func TestUpdatePing(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-001"})
	time.Sleep(10 * time.Millisecond)

	r.UpdatePing("tai-001")
	snap, _ := r.Get("tai-001")
	if snap.LastPing.Before(snap.ConnectedAt) {
		t.Error("LastPing should be after ConnectedAt")
	}
}

func TestUpdatePing_NonexistentNode(t *testing.T) {
	r := newTestRegistry()
	r.UpdatePing("ghost")
}

func TestGenerateChannelID_Unique(t *testing.T) {
	seen := make(map[string]bool)
	for i := 0; i < 100; i++ {
		id, err := generateChannelID()
		if err != nil {
			t.Fatalf("generateChannelID: %v", err)
		}
		if len(id) != 64 {
			t.Errorf("len = %d, want 64 hex chars", len(id))
		}
		if seen[id] {
			t.Fatalf("duplicate channel ID: %s", id)
		}
		seen[id] = true
	}
}

func TestConcurrentRegisterGet(t *testing.T) {
	r := newTestRegistry()
	var wg sync.WaitGroup

	for i := 0; i < 50; i++ {
		wg.Add(2)
		id := "tai-" + string(rune('A'+i%26))

		go func() {
			defer wg.Done()
			r.Register(&TaiNode{TaiID: id, Mode: "tunnel"})
		}()

		go func() {
			defer wg.Done()
			r.Get(id)
			r.List()
		}()
	}

	wg.Wait()
}

func TestOpenLocalListener_NodeNotFound(t *testing.T) {
	r := newTestRegistry()
	_, err := r.OpenLocalListener("ghost", 19100)
	if err == nil {
		t.Fatal("expected error for missing node")
	}
}

func TestRegister_SystemInfo(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{
		TaiID: "tai-001",
		System: types.SystemInfo{
			OS:       "linux",
			Arch:     "amd64",
			Hostname: "docker-host-01",
			NumCPU:   16,
		},
	})

	snap, ok := r.Get("tai-001")
	if !ok {
		t.Fatal("node not found")
	}
	if snap.System.OS != "linux" {
		t.Errorf("System.OS = %q, want linux", snap.System.OS)
	}
	if snap.System.Arch != "amd64" {
		t.Errorf("System.Arch = %q, want amd64", snap.System.Arch)
	}
	if snap.System.Hostname != "docker-host-01" {
		t.Errorf("System.Hostname = %q, want docker-host-01", snap.System.Hostname)
	}
	if snap.System.NumCPU != 16 {
		t.Errorf("System.NumCPU = %d, want 16", snap.System.NumCPU)
	}
}

func TestListByTeam(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-a", Auth: types.AuthInfo{TeamID: "team-dev"}})
	r.Register(&TaiNode{TaiID: "tai-b", Auth: types.AuthInfo{TeamID: "team-dev"}})
	r.Register(&TaiNode{TaiID: "tai-c", Auth: types.AuthInfo{TeamID: "team-ops"}})

	devNodes := r.ListByTeam("team-dev")
	if len(devNodes) != 2 {
		t.Errorf("ListByTeam(team-dev) = %d nodes, want 2", len(devNodes))
	}

	opsNodes := r.ListByTeam("team-ops")
	if len(opsNodes) != 1 {
		t.Errorf("ListByTeam(team-ops) = %d nodes, want 1", len(opsNodes))
	}

	empty := r.ListByTeam("team-ghost")
	if len(empty) != 0 {
		t.Errorf("ListByTeam(team-ghost) = %d nodes, want 0", len(empty))
	}
}

func TestStartHealthCheck_MarkOffline(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-direct", Mode: "direct"})
	r.Register(&TaiNode{TaiID: "tai-tunnel", Mode: "tunnel"})

	// Manually set LastPing to the past for the direct node.
	r.mu.Lock()
	r.nodes["tai-direct"].LastPing = time.Now().Add(-5 * time.Second)
	r.mu.Unlock()

	done := make(chan struct{})
	r.StartHealthCheck(done, 50*time.Millisecond, 2*time.Second, 10*time.Minute)
	defer close(done)

	time.Sleep(200 * time.Millisecond)

	snap, ok := r.Get("tai-direct")
	if !ok {
		t.Fatal("direct node should still exist")
	}
	if snap.Status != "offline" {
		t.Errorf("direct node Status = %q, want offline", snap.Status)
	}

	// Tunnel nodes should not be affected.
	snap2, ok := r.Get("tai-tunnel")
	if !ok {
		t.Fatal("tunnel node should still exist")
	}
	if snap2.Status != "online" {
		t.Errorf("tunnel node Status = %q, want online", snap2.Status)
	}
}

func TestStartHealthCheck_AutoCleanup(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-stale", Mode: "direct"})

	// Set LastPing far in the past so it exceeds both timeout and cleanupAfter.
	r.mu.Lock()
	r.nodes["tai-stale"].LastPing = time.Now().Add(-1 * time.Hour)
	r.mu.Unlock()

	done := make(chan struct{})
	r.StartHealthCheck(done, 50*time.Millisecond, 1*time.Second, 1*time.Second)
	defer close(done)

	time.Sleep(200 * time.Millisecond)

	if _, ok := r.Get("tai-stale"); ok {
		t.Error("stale node should have been auto-unregistered")
	}
}

func TestStartHealthCheck_PingKeepsAlive(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{TaiID: "tai-alive", Mode: "direct"})

	done := make(chan struct{})
	r.StartHealthCheck(done, 50*time.Millisecond, 2*time.Second, 10*time.Minute)
	defer close(done)

	// Continuously ping to keep the node alive.
	for i := 0; i < 4; i++ {
		time.Sleep(30 * time.Millisecond)
		r.UpdatePing("tai-alive")
	}

	snap, ok := r.Get("tai-alive")
	if !ok {
		t.Fatal("node should still exist")
	}
	if snap.Status != "online" {
		t.Errorf("Status = %q, want online", snap.Status)
	}
}

func TestNodeMeta_AuthInfo(t *testing.T) {
	r := newTestRegistry()
	r.Register(&TaiNode{
		TaiID: "tai-001",
		Auth: types.AuthInfo{
			Subject:  "user123",
			ClientID: "tai-001",
			Scope:    "tai:tunnel",
		},
	})

	snap, ok := r.Get("tai-001")
	if !ok {
		t.Fatal("node not found")
	}
	if snap.Auth.Subject != "user123" {
		t.Errorf("Auth.Subject = %q, want user123", snap.Auth.Subject)
	}
	if snap.Auth.Scope != "tai:tunnel" {
		t.Errorf("Auth.Scope = %q, want tai:tunnel", snap.Auth.Scope)
	}
}