iop/apps/control-plane/internal/wire/edge_server_test.go

package wire

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

	"go.uber.org/zap/zaptest"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	iop "iop/proto/gen/iop"
)

func startEdgeServer(t *testing.T) (*EdgeServer, int) {
	t.Helper()
	logger := zaptest.NewLogger(t)
	port := getFreePort(t)
	listenAddr := fmt.Sprintf("127.0.0.1:%d", port)

	server, err := NewEdgeServer(listenAddr, logger)
	if err != nil {
		t.Fatalf("NewEdgeServer failed: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	if err := server.Start(ctx); err != nil {
		t.Fatalf("edge server Start failed: %v", err)
	}
	t.Cleanup(func() { _ = server.Stop() })

	// Give a tiny moment for the listener to bind.
	time.Sleep(50 * time.Millisecond)
	return server, port
}

func dialEdge(t *testing.T, ctx context.Context, port int) *proto_socket.TcpClient {
	t.Helper()
	client, err := proto_socket.DialTcp(
		ctx,
		"127.0.0.1",
		port,
		proto_socket.DefaultHeartbeatIntervalSec,
		proto_socket.DefaultHeartbeatWaitSec,
		EdgeParserMap(),
	)
	if err != nil {
		t.Fatalf("failed to DialTcp: %v", err)
	}
	return client
}

func TestEdgeServerHelloRegistersEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeId:       "edge-dgx-group",
		EdgeName:     "DGX Group",
		Version:      "1.2.3",
		Capabilities: []string{"node-registry", "run-dispatch"},
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if !res.Accepted {
		t.Errorf("expected response.Accepted to be true, got false (reason=%q)", res.Reason)
	}
	if res.Protocol != Protocol {
		t.Errorf("expected response.Protocol to be %q, got %q", Protocol, res.Protocol)
	}
	if res.ServerTimeUnixNano <= 0 {
		t.Errorf("expected server time unix nano to be positive, got %d", res.ServerTimeUnixNano)
	}

	state, ok := server.Registry().Snapshot("edge-dgx-group")
	if !ok {
		t.Fatalf("edge not registered after hello")
	}
	if !state.Connected {
		t.Errorf("expected registry Connected to be true")
	}
	if state.EdgeName != "DGX Group" {
		t.Errorf("edge_name: got %q want %q", state.EdgeName, "DGX Group")
	}
	if len(state.Capabilities) != 2 {
		t.Errorf("capabilities: got %v", state.Capabilities)
	}
	if state.LastSeen.IsZero() {
		t.Errorf("expected LastSeen to be set")
	}
}

func TestEdgeServerRejectsMissingEdgeID(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeName: "no id edge",
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if res.Accepted {
		t.Errorf("expected response.Accepted to be false for missing edge_id")
	}
	if res.Reason == "" {
		t.Errorf("expected a rejection reason for missing edge_id")
	}
	if server.Registry().Len() != 0 {
		t.Errorf("expected registry to stay empty, got %d entries", server.Registry().Len())
	}
}

func TestEdgeServerDisconnectMarksEdgeDisconnected(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)

	req := &iop.EdgeHelloRequest{
		EdgeId:   "edge-dgx-group",
		EdgeName: "DGX Group",
	}
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if err := client.Close(); err != nil {
		t.Fatalf("client Close failed: %v", err)
	}

	// The server detects the closed connection asynchronously; poll until the
	// disconnect listener flips the registry state.
	deadline := time.Now().Add(2 * time.Second)
	for {
		state, ok := server.Registry().Snapshot("edge-dgx-group")
		if ok && !state.Connected {
			if state.DisconnectInfo.Reason == "" {
				t.Errorf("expected a disconnect reason to be recorded")
			}
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("edge was not marked disconnected in time (ok=%v)", ok)
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func sendEdgeHello(t *testing.T, client *proto_socket.TcpClient, edgeID string) {
	t.Helper()
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		&iop.EdgeHelloRequest{EdgeId: edgeID, EdgeName: "DGX Group"},
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}
}

func TestEdgeServerStaleDisconnectDoesNotClearReconnect(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const edgeID = "edge-dgx-group"

	// First connection enrolls the edge.
	conn1 := dialEdge(t, ctx, port)
	sendEdgeHello(t, conn1, edgeID)

	// Second connection with the same edge_id reconnects and becomes current.
	conn2 := dialEdge(t, ctx, port)
	defer conn2.Close()
	sendEdgeHello(t, conn2, edgeID)

	// The stale first connection closing must not clear the live reconnection.
	if err := conn1.Close(); err != nil {
		t.Fatalf("conn1 Close failed: %v", err)
	}

	// Give the server time to process the stale disconnect, then assert the edge
	// stays connected under the second connection's token.
	time.Sleep(300 * time.Millisecond)
	state, ok := server.Registry().Snapshot(edgeID)
	if !ok {
		t.Fatalf("edge missing from registry")
	}
	if !state.Connected {
		t.Fatalf("stale disconnect cleared the live reconnection; Connected=false")
	}
}

func TestEdgeHeartbeatWaitExceedsInterval(t *testing.T) {
	if EdgeHeartbeatWaitSec <= EdgeHeartbeatIntervalSec {
		t.Fatalf("EdgeHeartbeatWaitSec (%d) must exceed EdgeHeartbeatIntervalSec (%d)",
			EdgeHeartbeatWaitSec, EdgeHeartbeatIntervalSec)
	}
}