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

package wire

import (
	"testing"
	"time"

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

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

func TestEdgeParserMapEdgeHelloRoundTrip(t *testing.T) {
	parserMap := EdgeParserMap()
	typeName := proto_socket.TypeNameOf(&iop.EdgeHelloRequest{})
	parse, ok := parserMap[typeName]
	if !ok {
		t.Fatalf("missing parser for %s", typeName)
	}

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

	payload, err := proto.Marshal(req)
	if err != nil {
		t.Fatalf("marshal EdgeHelloRequest: %v", err)
	}

	got, err := parse(payload)
	if err != nil {
		t.Fatalf("parse EdgeHelloRequest: %v", err)
	}

	parsed, ok := got.(*iop.EdgeHelloRequest)
	if !ok {
		t.Fatalf("expected *iop.EdgeHelloRequest, got %T", got)
	}
	if parsed.EdgeId != req.EdgeId {
		t.Errorf("edge_id: got %q want %q", parsed.EdgeId, req.EdgeId)
	}
	if parsed.EdgeName != req.EdgeName {
		t.Errorf("edge_name: got %q want %q", parsed.EdgeName, req.EdgeName)
	}
	if parsed.Version != req.Version {
		t.Errorf("version: got %q want %q", parsed.Version, req.Version)
	}
	if len(parsed.Capabilities) != len(req.Capabilities) {
		t.Fatalf("capabilities length: got %d want %d", len(parsed.Capabilities), len(req.Capabilities))
	}
	if parsed.Capabilities[0] != req.Capabilities[0] || parsed.Capabilities[1] != req.Capabilities[1] {
		t.Errorf("capabilities: got %v want %v", parsed.Capabilities, req.Capabilities)
	}
	if parsed.Metadata["region"] != "local" {
		t.Errorf("metadata region: got %q want %q", parsed.Metadata["region"], "local")
	}
}

func TestEdgeParserMapEdgeHelloResponseRoundTrip(t *testing.T) {
	parserMap := EdgeParserMap()
	typeName := proto_socket.TypeNameOf(&iop.EdgeHelloResponse{})
	parse, ok := parserMap[typeName]
	if !ok {
		t.Fatalf("missing parser for %s", typeName)
	}

	res := &iop.EdgeHelloResponse{
		Accepted:           true,
		Protocol:           Protocol,
		ServerTimeUnixNano: time.Now().UnixNano(),
		Message:            "edge enrolled",
	}

	payload, err := proto.Marshal(res)
	if err != nil {
		t.Fatalf("marshal EdgeHelloResponse: %v", err)
	}

	got, err := parse(payload)
	if err != nil {
		t.Fatalf("parse EdgeHelloResponse: %v", err)
	}

	parsed, ok := got.(*iop.EdgeHelloResponse)
	if !ok {
		t.Fatalf("expected *iop.EdgeHelloResponse, got %T", got)
	}
	if !parsed.Accepted {
		t.Error("accepted: got false want true")
	}
	if parsed.Protocol != Protocol {
		t.Errorf("protocol: got %q want %q", parsed.Protocol, Protocol)
	}
	if parsed.ServerTimeUnixNano <= 0 {
		t.Errorf("server_time_unix_nano should be positive, got %d", parsed.ServerTimeUnixNano)
	}
}

func TestEdgeHelloContractDoesNotExposeNodeDirectScheduling(t *testing.T) {
	fields := (&iop.EdgeHelloRequest{}).ProtoReflect().Descriptor().Fields()
	forbidden := []protoreflect.Name{
		"node_id",
		"node_address",
		"schedule_request",
		"target",
	}
	for _, name := range forbidden {
		if fields.ByName(name) != nil {
			t.Fatalf("EdgeHelloRequest must not expose direct Node scheduling field %q", name)
		}
	}
}

func TestWireTransportBoundary(t *testing.T) {
	if Protocol != "protobuf-socket" {
		t.Fatalf("Protocol = %q, want %q", Protocol, "protobuf-socket")
	}
	if EdgeTransport != "proto-socket-tcp" {
		t.Fatalf("EdgeTransport = %q, want %q", EdgeTransport, "proto-socket-tcp")
	}
	if ClientTransport != "proto-socket-ws" {
		t.Fatalf("ClientTransport = %q, want %q", ClientTransport, "proto-socket-ws")
	}
}