iop/apps/node/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"

	"iop/apps/node/internal/transport"
	iop "iop/proto/gen/iop"
)

type mockHandler struct {
	runReqCh chan *iop.RunRequest
}

func (m *mockHandler) OnRunRequest(ctx context.Context, sess *transport.Session, req *iop.RunRequest) error {
	m.runReqCh <- req
	return sess.Send(&iop.RunEvent{RunId: req.GetRunId(), Type: "test_event"})
}

func (m *mockHandler) OnCapabilityRequest(ctx context.Context, sess *transport.Session) (*iop.CapabilityResponse, error) {
	return &iop.CapabilityResponse{
		NodeId: "test-node",
		Adapters: []*iop.AdapterInfo{
			{Name: "test-adapter", Models: []string{"v1"}, MaxConcurrency: 1},
		},
	}, nil
}

func (m *mockHandler) OnCancel(ctx context.Context, sess *transport.Session, req *iop.CancelRequest) error {
	return nil
}

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 edgeParserMap() toki.ParserMap {
	return toki.ParserMap{
		toki.TypeNameOf(&iop.RunEvent{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunEvent{}
			return m, proto.Unmarshal(b, m)
		},
		toki.TypeNameOf(&iop.CapabilityResponse{}): func(b []byte) (proto.Message, error) {
			m := &iop.CapabilityResponse{}
			return m, proto.Unmarshal(b, m)
		},
	}
}

func waitForAcceptedClient(t *testing.T, acceptedCh <-chan *toki.TcpClient) *toki.TcpClient {
	t.Helper()

	select {
	case client := <-acceptedCh:
		return client
	case <-time.After(2 * time.Second):
		t.Fatal("edge server did not accept connection")
		return nil
	}
}

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

	logger := zap.NewNop()
	listenAddr := getFreePort(t)
	host, portStr, _ := net.SplitHostPort(listenAddr)
	port := 0
	fmt.Sscanf(portStr, "%d", &port)

	// 1. Mock Edge 서버 구동
	acceptedCh := make(chan *toki.TcpClient, 1)
	server := toki.NewTcpServer(host, port, func(conn net.Conn) *toki.TcpClient {
		client := toki.NewTcpClient(conn, 30, 10, edgeParserMap())
		acceptedCh <- client
		return client
	})
	if err := server.Start(ctx); err != nil {
		t.Fatalf("failed to start mock edge server: %v", err)
	}
	defer server.Stop()

	// 2. Node 클라이언트 접속
	handler := &mockHandler{
		runReqCh: make(chan *iop.RunRequest, 1),
	}

	sess, err := transport.DialEdge(ctx, listenAddr, handler, "test-node", logger)
	if err != nil {
		t.Fatalf("failed to dial edge: %v", err)
	}
	defer sess.Close()

	edgeClient := waitForAcceptedClient(t, acceptedCh)
	runEventCh := make(chan *iop.RunEvent, 1)
	toki.AddListenerTyped[*iop.RunEvent](&edgeClient.Communicator, func(event *iop.RunEvent) {
		runEventCh <- event
	})

	// 3. Edge -> Node 로 CapabilityRequest 전송
	resp, err := toki.SendRequestTyped[*iop.CapabilityRequest, *iop.CapabilityResponse](
		&edgeClient.Communicator,
		&iop.CapabilityRequest{},
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("failed to send capability request: %v", err)
	}
	if resp.GetNodeId() != "test-node" {
		t.Fatalf("unexpected node id from capability response: %q", resp.GetNodeId())
	}

	// 4. Edge -> Node 로 RunRequest 전송
	runReq := &iop.RunRequest{
		RunId:   "test-run",
		Adapter: "test-adapter",
	}
	if err := edgeClient.Send(runReq); err != nil {
		t.Fatalf("failed to send run request: %v", err)
	}

	// 5. Node에서 RunRequest 수신 확인
	select {
	case receivedReq := <-handler.runReqCh:
		if receivedReq.GetRunId() != runReq.GetRunId() {
			t.Fatalf("expected run id %q, got %q", runReq.GetRunId(), receivedReq.GetRunId())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for run request on node handler")
	}

	select {
	case event := <-runEventCh:
		if event.GetRunId() != runReq.GetRunId() {
			t.Fatalf("expected run event id %q, got %q", runReq.GetRunId(), event.GetRunId())
		}
		if event.GetType() != "test_event" {
			t.Fatalf("expected run event type %q, got %q", "test_event", event.GetType())
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for run event from node session")
	}
}