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/common-proto-socket/go"

	edgenode "iop/apps/edge/internal/node"
	"iop/apps/edge/internal/transport"
	"iop/packages/config"
	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 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 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{
		{
			Alias: "test-node",
			Token: "test-token",
			Adapters: config.AdaptersConf{
				Ollama: config.OllamaConf{Enabled: true, BaseURL: "http://localhost:11434"},
			},
			Runtime: config.RuntimeConf{Concurrency: 2, WorkspaceRoot: "/tmp/iop/test"},
		},
	})
	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)
	}
	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 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()))
	}

	// 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)
	}
}