iop/apps/node/internal/transport/parser_test.go

package transport

import (
	"testing"

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

	iop "iop/proto/gen/iop"
)

func TestNodeParserMap_RunRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.RunRequest{
		RunId:       "run-1",
		Adapter:     "mock",
		Target:      "v1",
		SessionId:   "session-a",
		Background:  true,
		SessionMode: iop.RunSessionMode_RUN_SESSION_MODE_REQUIRE_EXISTING,
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.RunRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetBackground() != original.GetBackground() ||
		got.GetSessionMode() != original.GetSessionMode() {
		t.Fatalf("unexpected run request: %+v", got)
	}
}

func TestNodeParserMap_ProviderTunnelRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.ProviderTunnelRequest{
		RunId:      "run-1",
		TunnelId:   "tunnel-1",
		Adapter:    "openai_compat",
		Target:     "qwen",
		Method:     "POST",
		Path:       "/v1/chat/completions",
		Headers:    map[string]string{"content-type": "application/json"},
		Body:       []byte(`{"stream":true}`),
		Stream:     true,
		TimeoutSec: 30,
		Metadata:   map[string]string{"openai_model": "qwen"},
		SessionId:  "openai",
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.ProviderTunnelRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetTunnelId() != original.GetTunnelId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetMethod() != original.GetMethod() ||
		got.GetPath() != original.GetPath() ||
		got.GetHeaders()["content-type"] != "application/json" ||
		string(got.GetBody()) != string(original.GetBody()) ||
		got.GetStream() != original.GetStream() ||
		got.GetTimeoutSec() != original.GetTimeoutSec() ||
		got.GetMetadata()["openai_model"] != "qwen" ||
		got.GetSessionId() != original.GetSessionId() {
		t.Fatalf("unexpected tunnel request: %+v", got)
	}
}

func TestNodeParserMap_CancelRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.CancelRequest{
		RunId:     "run-1",
		Adapter:   "cli",
		Target:    "codex",
		SessionId: "session-a",
		Action:    iop.CancelAction_CANCEL_ACTION_TERMINATE_SESSION,
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.CancelRequest)
	if got.GetRunId() != original.GetRunId() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() ||
		got.GetAction() != original.GetAction() {
		t.Fatalf("unexpected cancel request: %+v", got)
	}
}

func TestNodeParserMap_EdgeNodeEvent(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.EdgeNodeEvent{
		Type:   "edge.disconnected",
		Source: "node",
		NodeId: "node-1",
		Alias:  "local-node",
		Reason: "transport_closed",
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.EdgeNodeEvent)
	if got.GetType() != original.GetType() ||
		got.GetSource() != original.GetSource() ||
		got.GetNodeId() != original.GetNodeId() ||
		got.GetAlias() != original.GetAlias() ||
		got.GetReason() != original.GetReason() {
		t.Fatalf("unexpected edge node event: %+v", got)
	}
}

func TestNodeParserMap_NodeCommandRequest(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.NodeCommandRequest{
		Type:      iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
		Adapter:   "cli",
		Target:    "codex",
		SessionId: "default",
	}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got := parsed.(*iop.NodeCommandRequest)
	if got.GetType() != original.GetType() ||
		got.GetAdapter() != original.GetAdapter() ||
		got.GetTarget() != original.GetTarget() ||
		got.GetSessionId() != original.GetSessionId() {
		t.Fatalf("unexpected node command request: %+v", got)
	}
}

func TestNodeParserMap_NodeCommandRequest_NewTypes(t *testing.T) {
	parsers := nodeParserMap()
	types := []iop.NodeCommandType{
		iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
		iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST,
		iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS,
	}
	for _, cmdType := range types {
		original := &iop.NodeCommandRequest{
			RequestId: "req-1",
			Type:      cmdType,
			Adapter:   "cli",
			Target:    "codex",
			SessionId: "default",
		}
		payload, err := proto.Marshal(original)
		if err != nil {
			t.Fatalf("marshal %v: %v", cmdType, err)
		}
		parser, ok := parsers[toki.TypeNameOf(original)]
		if !ok {
			t.Fatalf("parser not found for NodeCommandRequest")
		}
		parsed, err := parser(payload)
		if err != nil {
			t.Fatalf("parse %v: %v", cmdType, err)
		}
		got := parsed.(*iop.NodeCommandRequest)
		if got.GetType() != cmdType {
			t.Fatalf("type: got %v want %v", got.GetType(), cmdType)
		}
	}
}

func TestNodeParserMap_RegisterResponse(t *testing.T) {
	parsers := nodeParserMap()
	original := &iop.RegisterResponse{Accepted: true, NodeId: "node-1", Alias: "local-node"}
	payload, err := proto.Marshal(original)
	if err != nil {
		t.Fatalf("marshal: %v", err)
	}

	key := toki.TypeNameOf(original)
	parser, ok := parsers[key]
	if !ok {
		t.Fatalf("parser not found for key: %s", key)
	}

	parsed, err := parser(payload)
	if err != nil {
		t.Fatalf("parse: %v", err)
	}
	got, ok := parsed.(*iop.RegisterResponse)
	if !ok {
		t.Fatalf("unexpected register response type: %T", parsed)
	}
	if !got.GetAccepted() || got.GetNodeId() != original.GetNodeId() || got.GetAlias() != original.GetAlias() {
		t.Fatalf("unexpected register response: %+v", got)
	}
}