iop/apps/edge/internal/transport/integration_test.go

package transport_test

import (
	"context"
	"fmt"
	"net"
	"testing"
	"time"

	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"

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

	edgenode "iop/apps/edge/internal/node"
	"iop/apps/edge/internal/transport"
	"iop/packages/go/config"
	eventpkg "iop/packages/go/events"
	iop "iop/proto/gen/iop"
)

func getFreePort(t *testing.T) string {
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := l.Addr().String()
	l.Close()
	return addr
}

func dialNode(t *testing.T, ctx context.Context, listenAddr string) *toki.TcpClient {
	t.Helper()
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port := 0
	fmt.Sscanf(portStr, "%d", &port)
	nodeParser := toki.ParserMap{
		toki.TypeNameOf(&iop.RegisterResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
	client, err := toki.DialTcp(ctx, host, port, 30, 10, nodeParser)
	if err != nil {
		t.Fatalf("dial edge server: %v", err)
	}
	return client
}

func waitForRegistryEntry(ctx context.Context, registry *edgenode.Registry, nodeID string) (*edgenode.NodeEntry, bool) {
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()

	for {
		if entry, ok := registry.Get(nodeID); ok {
			return entry, true
		}

		select {
		case <-ctx.Done():
			return nil, false
		case <-ticker.C:
		}
	}
}

func waitForRegistryAbsent(ctx context.Context, registry *edgenode.Registry, nodeID string) bool {
	ticker := time.NewTicker(25 * time.Millisecond)
	defer ticker.Stop()

	for {
		if _, ok := registry.Get(nodeID); !ok {
			return true
		}

		select {
		case <-ctx.Done():
			return false
		case <-ticker.C:
		}
	}
}

func TestEdgeServerGenericRegistrationKind(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-test-01",
			Alias:     "generic-node",
			Token:     "node-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	client := dialNode(t, ctx, listenAddr)
	defer client.Close()

	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: "node-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("register request: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted register response, got reason %q", resp.GetReason())
	}

	entry, ok := waitForRegistryEntry(ctx, registry, "node-test-01")
	if !ok {
		t.Fatal("node was not registered in edge registry within timeout")
	}
	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 TestEdgeServerDuplicateRegistrationReason(t *testing.T) {
	cases := []struct {
		name       string
		kind       string
		wantReason string
	}{
		{"generic", config.AgentKindGenericNode, "node already connected"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			logger := zap.NewNop()
			listenAddr := getFreePort(t)
			registry := edgenode.NewRegistry()
			nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
				{ID: "dup-01", Alias: "dup", Token: "dup-token", AgentKind: tc.kind},
			})
			if err != nil {
				t.Fatalf("load node store: %v", err)
			}

			server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
			if err != nil {
				t.Fatalf("new server: %v", err)
			}
			if err := server.Start(ctx); err != nil {
				t.Fatalf("start server: %v", err)
			}
			defer server.Stop()

			first := dialNode(t, ctx, listenAddr)
			defer first.Close()
			resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
				&first.Communicator, &iop.RegisterRequest{Token: "dup-token"}, 2*time.Second)
			if err != nil {
				t.Fatalf("first register: %v", err)
			}
			if !resp1.GetAccepted() {
				t.Fatalf("expected first registration accepted, got reason %q", resp1.GetReason())
			}
			if _, ok := waitForRegistryEntry(ctx, registry, "dup-01"); !ok {
				t.Fatal("first registration not visible in registry within timeout")
			}

			second := dialNode(t, ctx, listenAddr)
			defer second.Close()
			resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
				&second.Communicator, &iop.RegisterRequest{Token: "dup-token"}, 2*time.Second)
			if err != nil {
				t.Fatalf("second register: %v", err)
			}
			if resp2.GetAccepted() {
				t.Fatal("expected second registration to be rejected")
			}
			if resp2.GetReason() != tc.wantReason {
				t.Fatalf("duplicate reason: got %q want %q", resp2.GetReason(), tc.wantReason)
			}
		})
	}
}

func TestEdgeServerReconnectAfterUnregisterAccepted(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{ID: "reconnect-01", Alias: "reconnect", Token: "reconnect-token", AgentKind: config.AgentKindGenericNode},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	first := dialNode(t, ctx, listenAddr)
	resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&first.Communicator, &iop.RegisterRequest{Token: "reconnect-token"}, 2*time.Second)
	if err != nil {
		t.Fatalf("first register: %v", err)
	}
	if !resp1.GetAccepted() {
		t.Fatalf("expected first registration accepted, got reason %q", resp1.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("first registration not visible in registry within timeout")
	}

	if err := first.Close(); err != nil {
		t.Fatalf("close first client: %v", err)
	}
	if ok := waitForRegistryAbsent(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("registry still contains node after disconnect")
	}

	second := dialNode(t, ctx, listenAddr)
	defer second.Close()
	resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&second.Communicator, &iop.RegisterRequest{Token: "reconnect-token"}, 2*time.Second)
	if err != nil {
		t.Fatalf("second register: %v", err)
	}
	if !resp2.GetAccepted() {
		t.Fatalf("expected reconnect after unregister accepted, got reason %q", resp2.GetReason())
	}
	if _, ok := waitForRegistryEntry(ctx, registry, "reconnect-01"); !ok {
		t.Fatal("reconnect registration not visible in registry within timeout")
	}
}

func TestEdgeServerIntegration(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:    "node-test-01",
			Alias: "test-node",
			Token: "test-token",
			Adapters: config.AdaptersConf{
				Ollama: config.OllamaConf{Enabled: true, BaseURL: "http://localhost:11434"},
			},
			Runtime: config.RuntimeConf{Concurrency: 2},
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	// 1. Edge 서버 구동
	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("failed to create server: %v", err)
	}
	nodeEventCh := make(chan *iop.EdgeNodeEvent, 4)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("failed to start server: %v", err)
	}
	defer server.Stop()

	// 2. Mock Node 클라이언트 접속
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port := 0
	fmt.Sscanf(portStr, "%d", &port)

	// Node parser map (mock)
	nodeParser := toki.ParserMap{
		toki.TypeNameOf(&iop.RegisterResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.RegisterResponse{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	client, err := toki.DialTcp(ctx, host, port, 30, 10, nodeParser)
	if err != nil {
		t.Fatalf("failed to dial edge server: %v", err)
	}
	defer client.Close()

	// 3. Node가 RegisterRequest를 보내고 RegisterResponse를 받는다
	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: "test-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("register request: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted register response, got reason %q", resp.GetReason())
	}
	rec, ok := nodeStore.FindByToken("test-token")
	if !ok {
		t.Fatal("expected record for test-token in store")
	}
	wantNodeID := rec.ID
	if wantNodeID != "node-test-01" {
		t.Fatalf("expected stored id to match explicit config id, got %q", wantNodeID)
	}
	if resp.GetNodeId() != wantNodeID {
		t.Fatalf("expected node id %q, got %q", wantNodeID, resp.GetNodeId())
	}
	if resp.GetAlias() != "test-node" {
		t.Fatalf("expected alias %q, got %q", "test-node", resp.GetAlias())
	}
	if resp.GetConfig().GetRuntime().GetConcurrency() != 2 {
		t.Fatalf("expected runtime concurrency 2, got %d", resp.GetConfig().GetRuntime().GetConcurrency())
	}
	if len(resp.GetConfig().GetAdapters()) != 2 {
		t.Fatalf("expected mock and ollama adapters, got %d", len(resp.GetConfig().GetAdapters()))
	}
	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeConnected {
			t.Fatalf("event type: got %q want %q", event.GetType(), eventpkg.TypeNodeConnected)
		}
		if event.GetNodeId() != wantNodeID || event.GetAlias() != "test-node" || event.GetReason() != eventpkg.ReasonRegistered {
			t.Fatalf("unexpected connected event: %+v", event)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for node connected event")
	}

	// 4. Registry 등록 여부 확인
	entry, ok := waitForRegistryEntry(ctx, registry, wantNodeID)
	if !ok {
		t.Fatal("node was not registered in edge registry within timeout")
	}

	if got := registry.Count(); got != 1 {
		t.Fatalf("expected registry count 1, got %d", got)
	}
	if entry == nil {
		t.Fatal("expected registry entry, got nil")
	}
	if got := entry.NodeID; got != wantNodeID {
		t.Fatalf("expected node id %q, got %q", wantNodeID, got)
	}
	if got := entry.Alias; got != "test-node" {
		t.Fatalf("expected alias %q, got %q", "test-node", got)
	}

	// 5. Alias로 Resolve 확인
	if e, err := registry.Resolve("test-node"); err != nil || e.NodeID != wantNodeID {
		t.Fatalf("failed to resolve by alias %q: %v", "test-node", err)
	}

	if err := client.Close(); err != nil {
		t.Fatalf("close client: %v", err)
	}
	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeDisconnected {
			t.Fatalf("event type: got %q want %q", event.GetType(), eventpkg.TypeNodeDisconnected)
		}
		if event.GetNodeId() != wantNodeID || event.GetAlias() != "test-node" || event.GetReason() != eventpkg.ReasonTransportClosed {
			t.Fatalf("unexpected disconnected event: %+v", event)
		}
		if event.GetMetadata()[eventpkg.MetadataTransportCloseReason] != toki.DisconnectReasonRemoteClosed {
			t.Fatalf("transport close reason: got %q want %q", event.GetMetadata()[eventpkg.MetadataTransportCloseReason], toki.DisconnectReasonRemoteClosed)
		}
		if event.GetMetadata()[eventpkg.MetadataTransportCloseError] == "" {
			t.Fatalf("expected transport close error metadata, got %+v", event.GetMetadata())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for node disconnected event")
	}
}

func TestEdgeServerRegistrationFailureReasons(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-fail-01",
			Alias:     "fail-node",
			Token:     "valid-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}

	nodeEventCh := make(chan *iop.EdgeNodeEvent, 10)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	// 1. Unknown Token Test
	client1 := dialNode(t, ctx, listenAddr)
	resp1, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client1.Communicator,
		&iop.RegisterRequest{Token: "unknown-token-12345"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 1: %v", err)
	}
	if resp1.GetAccepted() {
		t.Fatal("expected rejected")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeRegistrationFailed {
			t.Fatalf("expected registration failed type, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonUnknownToken {
			t.Fatalf("expected reason unknown_token, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonUnknownToken {
			t.Fatalf("expected metadata failure reason, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataTokenPrefix] != "unknown-..." {
			t.Fatalf("expected token prefix unknown-..., got %q", event.GetMetadata()[eventpkg.MetadataTokenPrefix])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for unknown token registration failed event")
	}
	client1.Close()

	// 2. Duplicate Connection Test
	client2 := dialNode(t, ctx, listenAddr)
	resp2, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client2.Communicator,
		&iop.RegisterRequest{Token: "valid-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 2: %v", err)
	}
	if !resp2.GetAccepted() {
		t.Fatalf("expected register 2 accepted, got %q", resp2.GetReason())
	}
	// Consume connection event
	<-nodeEventCh

	client3 := dialNode(t, ctx, listenAddr)
	resp3, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client3.Communicator,
		&iop.RegisterRequest{Token: "valid-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("send register 3: %v", err)
	}
	if resp3.GetAccepted() {
		t.Fatal("expected duplicate connection registration rejected")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeRegistrationFailed {
			t.Fatalf("expected registration failed type, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonDuplicateConnection {
			t.Fatalf("expected reason duplicate_connection, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonDuplicateConnection {
			t.Fatalf("expected failure reason metadata, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataAgentKind] != config.AgentKindGenericNode {
			t.Fatalf("expected agent kind metadata, got %q", event.GetMetadata()[eventpkg.MetadataAgentKind])
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for duplicate registration failed event")
	}
	client2.Close()
	client3.Close()
}

func TestEdgeServerHeartbeatTimeoutDisconnect(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	registry := edgenode.NewRegistry()
	nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
		{
			ID:        "node-timeout-01",
			Alias:     "timeout-node",
			Token:     "timeout-token",
			AgentKind: config.AgentKindGenericNode,
		},
	})
	if err != nil {
		t.Fatalf("load node store: %v", err)
	}

	server, err := transport.NewServer(listenAddr, registry, nodeStore, logger)
	if err != nil {
		t.Fatalf("new server: %v", err)
	}
	server.HeartbeatInterval = 1
	server.HeartbeatWait = 1

	nodeEventCh := make(chan *iop.EdgeNodeEvent, 10)
	server.SetNodeEventHandler(func(event *iop.EdgeNodeEvent) {
		nodeEventCh <- event
	})

	if err := server.Start(ctx); err != nil {
		t.Fatalf("start server: %v", err)
	}
	defer server.Stop()

	client := dialNode(t, ctx, listenAddr)
	defer client.Close()
	client.RemoveListeners(toki.TypeNameOf(&packets.HeartBeat{}))

	resp, err := toki.SendRequestTyped[*iop.RegisterRequest, *iop.RegisterResponse](
		&client.Communicator,
		&iop.RegisterRequest{Token: "timeout-token"},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("register request: %v", err)
	}
	if !resp.GetAccepted() {
		t.Fatalf("expected accepted, got %q", resp.GetReason())
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeConnected {
			t.Fatalf("expected connected event, got %q", event.GetType())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for connected event")
	}

	select {
	case event := <-nodeEventCh:
		if event.GetType() != eventpkg.TypeNodeDisconnected {
			t.Fatalf("expected disconnected event, got %q", event.GetType())
		}
		if event.GetReason() != eventpkg.ReasonHeartbeatTimeout {
			t.Fatalf("expected reason heartbeat_timeout, got %q", event.GetReason())
		}
		if event.GetMetadata()[eventpkg.MetadataFailureReason] != eventpkg.ReasonHeartbeatTimeout {
			t.Fatalf("expected failure reason heartbeat_timeout, got %q", event.GetMetadata()[eventpkg.MetadataFailureReason])
		}
		if event.GetMetadata()[eventpkg.MetadataLifecycleState] != edgenode.LifecycleFailed {
			t.Fatalf("expected lifecycle state failed, got %q", event.GetMetadata()[eventpkg.MetadataLifecycleState])
		}
	case <-time.After(3 * time.Second):
		t.Fatal("timeout waiting for heartbeat timeout disconnected event")
	}
}