toki
cc87db264b
- Add mock resource activation for G06 review followup alignment - Expand edgevalidate tests and refactor validation logic - Update node mapper with improved handling - Add integration and server test fixes - Update Go config with new options - Fix E2E scripts (cli-workspace, lemonade, ollama, vllm)
449 lines
13 KiB
Go
449 lines
13 KiB
Go
package transport
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
toki "git.toki-labs.com/toki/proto-socket/go"
|
|
"go.uber.org/zap"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
edgenode "iop/apps/edge/internal/node"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func TestEdgeParserMap_NodeCommandResponse(t *testing.T) {
|
|
parsers := edgeParserMap()
|
|
original := &iop.NodeCommandResponse{
|
|
Type: iop.NodeCommandType_NODE_COMMAND_TYPE_USAGE_STATUS,
|
|
Adapter: "cli",
|
|
Target: "codex",
|
|
SessionId: "default",
|
|
UsageStatus: &iop.AgentUsageStatus{
|
|
RawOutput: "test",
|
|
},
|
|
}
|
|
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.NodeCommandResponse)
|
|
if got.GetType() != original.GetType() ||
|
|
got.GetAdapter() != original.GetAdapter() ||
|
|
got.GetTarget() != original.GetTarget() ||
|
|
got.GetSessionId() != original.GetSessionId() ||
|
|
got.GetUsageStatus().GetRawOutput() != original.GetUsageStatus().GetRawOutput() {
|
|
t.Fatalf("unexpected node command response: %+v", got)
|
|
}
|
|
}
|
|
|
|
func TestEdgeParserMap_NodeCommandResponse_NewTypes(t *testing.T) {
|
|
parsers := edgeParserMap()
|
|
cases := []struct {
|
|
cmdType iop.NodeCommandType
|
|
result map[string]string
|
|
providerSnapshots []*iop.ProviderSnapshot
|
|
}{
|
|
{
|
|
cmdType: iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES,
|
|
result: map[string]string{"targets": "codex,antigravity"},
|
|
providerSnapshots: []*iop.ProviderSnapshot{
|
|
{
|
|
Adapter: "cli",
|
|
Status: "available",
|
|
Capacity: 4,
|
|
InFlight: 2,
|
|
Queued: 0,
|
|
},
|
|
},
|
|
},
|
|
{iop.NodeCommandType_NODE_COMMAND_TYPE_SESSION_LIST, map[string]string{"count": "2"}, nil},
|
|
{iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, map[string]string{"connected": "true"}, nil},
|
|
}
|
|
for _, tc := range cases {
|
|
original := &iop.NodeCommandResponse{
|
|
RequestId: "req-1",
|
|
Type: tc.cmdType,
|
|
Adapter: "cli",
|
|
Target: "codex",
|
|
SessionId: "default",
|
|
Result: tc.result,
|
|
ProviderSnapshots: tc.providerSnapshots,
|
|
}
|
|
payload, err := proto.Marshal(original)
|
|
if err != nil {
|
|
t.Fatalf("marshal %v: %v", tc.cmdType, err)
|
|
}
|
|
parser, ok := parsers[toki.TypeNameOf(original)]
|
|
if !ok {
|
|
t.Fatalf("parser not found for NodeCommandResponse")
|
|
}
|
|
parsed, err := parser(payload)
|
|
if err != nil {
|
|
t.Fatalf("parse %v: %v", tc.cmdType, err)
|
|
}
|
|
got := parsed.(*iop.NodeCommandResponse)
|
|
if got.GetType() != tc.cmdType {
|
|
t.Fatalf("type: got %v want %v", got.GetType(), tc.cmdType)
|
|
}
|
|
for k, v := range tc.result {
|
|
if got.GetResult()[k] != v {
|
|
t.Fatalf("result[%q]: got %q want %q", k, got.GetResult()[k], v)
|
|
}
|
|
}
|
|
if len(got.GetProviderSnapshots()) != len(tc.providerSnapshots) {
|
|
t.Fatalf("provider snapshots count: got %d want %d", len(got.GetProviderSnapshots()), len(tc.providerSnapshots))
|
|
}
|
|
for i, snap := range tc.providerSnapshots {
|
|
g := got.GetProviderSnapshots()[i]
|
|
if g.GetAdapter() != snap.GetAdapter() ||
|
|
g.GetStatus() != snap.GetStatus() ||
|
|
g.GetCapacity() != snap.GetCapacity() ||
|
|
g.GetInFlight() != snap.GetInFlight() ||
|
|
g.GetQueued() != snap.GetQueued() {
|
|
t.Fatalf("provider snapshot mismatch at index %d: %+v vs %+v", i, g, snap)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestEdgeParserMap_EdgeNodeEvent(t *testing.T) {
|
|
parsers := edgeParserMap()
|
|
original := &iop.EdgeNodeEvent{
|
|
Type: "node.disconnected",
|
|
Source: "edge",
|
|
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 TestServerEnrichesRunEventNodeAlias(t *testing.T) {
|
|
registry := edgenode.NewRegistry()
|
|
registry.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alias-1"})
|
|
s := &Server{registry: registry}
|
|
|
|
event := &iop.RunEvent{RunId: "run-1", Type: "start", NodeId: "node-1"}
|
|
s.enrichRunEvent(event)
|
|
if event.GetNodeAlias() != "alias-1" {
|
|
t.Fatalf("expected alias-1, got %q", event.GetNodeAlias())
|
|
}
|
|
|
|
preset := &iop.RunEvent{RunId: "run-2", Type: "start", NodeId: "node-1", NodeAlias: "preset"}
|
|
s.enrichRunEvent(preset)
|
|
if preset.GetNodeAlias() != "preset" {
|
|
t.Fatalf("expected preset to be preserved, got %q", preset.GetNodeAlias())
|
|
}
|
|
|
|
unknown := &iop.RunEvent{RunId: "run-3", Type: "start", NodeId: "node-x"}
|
|
s.enrichRunEvent(unknown)
|
|
if unknown.GetNodeAlias() != "" {
|
|
t.Fatalf("expected empty alias for unknown node, got %q", unknown.GetNodeAlias())
|
|
}
|
|
}
|
|
|
|
// TestEdgeParserMap_NodeConfigRefreshResponse verifies the parser map includes
|
|
// the NodeConfigRefreshResponse type so Edge can receive Node ack messages.
|
|
func TestEdgeParserMap_NodeConfigRefreshResponse(t *testing.T) {
|
|
parsers := edgeParserMap()
|
|
original := &iop.NodeConfigRefreshResponse{
|
|
RequestId: "req-refresh-1",
|
|
Status: iop.NodeConfigRefreshStatus_NODE_CONFIG_REFRESH_STATUS_APPLIED,
|
|
}
|
|
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.NodeConfigRefreshResponse)
|
|
if got.GetRequestId() != original.GetRequestId() || got.GetStatus() != original.GetStatus() {
|
|
t.Fatalf("unexpected response: %+v", got)
|
|
}
|
|
}
|
|
|
|
// TestServerPushConfigRefreshSkippedWhenNodeDisconnected verifies that a node
|
|
// configured in NodeStore but absent from the live registry is recorded as
|
|
// skipped. This represents the real disconnected-node path: the disconnect
|
|
// listener removes the entry from the registry, so configured-but-disconnected
|
|
// nodes never appear in registry.All().
|
|
func TestServerPushConfigRefreshSkippedWhenNodeDisconnected(t *testing.T) {
|
|
nodeStore, err := edgenode.LoadFromConfig([]config.NodeDefinition{
|
|
{ID: "node-1", Alias: "alias-1", Token: "tok-1"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("load node store: %v", err)
|
|
}
|
|
// Registry has no entry for node-1 (simulates post-disconnect state).
|
|
registry := edgenode.NewRegistry()
|
|
|
|
s := &Server{
|
|
registry: registry,
|
|
nodeStore: nodeStore,
|
|
logger: zap.NewNop(),
|
|
HeartbeatInterval: 30,
|
|
HeartbeatWait: 45,
|
|
}
|
|
|
|
req := &iop.NodeConfigRefreshRequest{RequestId: "req-1"}
|
|
results := s.PushConfigRefresh(context.Background(), req)
|
|
|
|
if len(results) != 1 {
|
|
t.Fatalf("expected 1 result, got %d", len(results))
|
|
}
|
|
if results[0].NodeID != "node-1" {
|
|
t.Fatalf("unexpected node id: %q", results[0].NodeID)
|
|
}
|
|
if results[0].Status != "skipped" {
|
|
t.Fatalf("expected status=skipped for absent registry entry, got %q", results[0].Status)
|
|
}
|
|
}
|
|
|
|
// TestServerPushConfigRefreshEmptyNodeStoreProducesNoResults verifies that
|
|
// PushConfigRefresh with an empty NodeStore returns an empty result slice.
|
|
func TestServerPushConfigRefreshEmptyNodeStoreProducesNoResults(t *testing.T) {
|
|
nodeStore, err := edgenode.LoadFromConfig(nil)
|
|
if err != nil {
|
|
t.Fatalf("load node store: %v", err)
|
|
}
|
|
s := &Server{
|
|
registry: edgenode.NewRegistry(),
|
|
nodeStore: nodeStore,
|
|
logger: zap.NewNop(),
|
|
HeartbeatInterval: 30,
|
|
HeartbeatWait: 45,
|
|
}
|
|
results := s.PushConfigRefresh(context.Background(), &iop.NodeConfigRefreshRequest{RequestId: "req-empty"})
|
|
if len(results) != 0 {
|
|
t.Fatalf("expected 0 results for empty node store, got %d", len(results))
|
|
}
|
|
}
|
|
|
|
func findCLIAdapter(payload *iop.NodeConfigPayload) *iop.AdapterConfig {
|
|
for _, a := range payload.GetAdapters() {
|
|
if a.GetType() == "cli" {
|
|
return a
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func TestBuildConfigPayload_CLIOneof(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-local",
|
|
Alias: "local",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{
|
|
Enabled: true,
|
|
Profiles: map[string]config.CLIProfileConf{
|
|
"codex": {
|
|
Command: "codex",
|
|
Args: []string{"--foo", "--bar"},
|
|
Env: []string{"KEY=val"},
|
|
Persistent: true,
|
|
Terminal: true,
|
|
ResponseIdleTimeoutMS: 1500,
|
|
StartupIdleTimeoutMS: 300,
|
|
OutputFormat: "codex-json",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("build config payload: %v", err)
|
|
}
|
|
|
|
cli := findCLIAdapter(payload)
|
|
if cli == nil {
|
|
t.Fatal("expected cli adapter")
|
|
}
|
|
if cli.GetSettings() != nil {
|
|
t.Fatalf("expected settings to be nil for cli adapter, got %v", cli.GetSettings())
|
|
}
|
|
cliCfg := cli.GetCli()
|
|
if cliCfg == nil {
|
|
t.Fatal("expected typed cli config")
|
|
}
|
|
prof, ok := cliCfg.GetProfiles()["codex"]
|
|
if !ok {
|
|
t.Fatal("expected codex profile")
|
|
}
|
|
if prof.GetCommand() != "codex" {
|
|
t.Fatalf("command: %q", prof.GetCommand())
|
|
}
|
|
if got := prof.GetArgs(); len(got) != 2 || got[0] != "--foo" || got[1] != "--bar" {
|
|
t.Fatalf("args: %#v", got)
|
|
}
|
|
if got := prof.GetEnv(); len(got) != 1 || got[0] != "KEY=val" {
|
|
t.Fatalf("env: %#v", got)
|
|
}
|
|
if !prof.GetPersistent() {
|
|
t.Fatal("persistent")
|
|
}
|
|
if !prof.GetTerminal() {
|
|
t.Fatal("terminal")
|
|
}
|
|
if prof.GetResponseIdleTimeoutMs() != 1500 {
|
|
t.Fatalf("response_idle_timeout_ms: %d", prof.GetResponseIdleTimeoutMs())
|
|
}
|
|
if prof.GetStartupIdleTimeoutMs() != 300 {
|
|
t.Fatalf("startup_idle_timeout_ms: %d", prof.GetStartupIdleTimeoutMs())
|
|
}
|
|
if prof.GetOutputFormat() != "codex-json" {
|
|
t.Fatalf("output_format: %q", prof.GetOutputFormat())
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_CLICompletionMarker(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
ID: "node-local",
|
|
Alias: "local",
|
|
Adapters: config.AdaptersConf{
|
|
CLI: config.CLIConf{
|
|
Enabled: true,
|
|
Profiles: map[string]config.CLIProfileConf{
|
|
"codex": {
|
|
Command: "codex",
|
|
Persistent: true,
|
|
CompletionMarker: config.CompletionMarkerConf{
|
|
Line: "[[DONE]]",
|
|
Regex: `^\[\[DONE\]\]$`,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("build config payload: %v", err)
|
|
}
|
|
cli := findCLIAdapter(payload)
|
|
if cli == nil {
|
|
t.Fatal("expected cli adapter")
|
|
}
|
|
prof := cli.GetCli().GetProfiles()["codex"]
|
|
marker := prof.GetCompletionMarker()
|
|
if marker == nil {
|
|
t.Fatal("expected completion_marker")
|
|
}
|
|
if marker.GetLine() != "[[DONE]]" {
|
|
t.Fatalf("line: %q", marker.GetLine())
|
|
}
|
|
if marker.GetRegex() != `^\[\[DONE\]\]$` {
|
|
t.Fatalf("regex: %q", marker.GetRegex())
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_OllamaVllmOneof(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama", Enabled: true, BaseURL: "http://localhost:11434"},
|
|
},
|
|
VllmInstances: []config.VllmInstanceConf{
|
|
{Name: "vllm", Enabled: true, Endpoint: "http://localhost:8000"},
|
|
},
|
|
},
|
|
}
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("build: %v", err)
|
|
}
|
|
var ollama, vllm *iop.AdapterConfig
|
|
for _, a := range payload.GetAdapters() {
|
|
switch a.GetType() {
|
|
case "ollama":
|
|
ollama = a
|
|
case "vllm":
|
|
vllm = a
|
|
}
|
|
}
|
|
if ollama == nil || ollama.GetOllama().GetBaseUrl() != "http://localhost:11434" {
|
|
t.Fatalf("ollama: %+v", ollama)
|
|
}
|
|
if vllm == nil || vllm.GetVllm().GetEndpoint() != "http://localhost:8000" {
|
|
t.Fatalf("vllm: %+v", vllm)
|
|
}
|
|
}
|
|
|
|
func TestBuildConfigPayload_AllAdaptersSettingsNil(t *testing.T) {
|
|
rec := &edgenode.NodeRecord{
|
|
Adapters: config.AdaptersConf{
|
|
Ollama: config.OllamaConf{Enabled: true, BaseURL: "http://localhost:11434"},
|
|
Vllm: config.VllmConf{Enabled: true, Endpoint: "http://localhost:8000"},
|
|
Mock: config.MockConf{Enabled: true},
|
|
CLI: config.CLIConf{
|
|
Enabled: true,
|
|
Profiles: map[string]config.CLIProfileConf{"default": {Command: "echo"}},
|
|
},
|
|
},
|
|
}
|
|
payload, err := edgenode.BuildConfigPayload(rec)
|
|
if err != nil {
|
|
t.Fatalf("build: %v", err)
|
|
}
|
|
for _, a := range payload.GetAdapters() {
|
|
if a.GetSettings() != nil {
|
|
t.Fatalf("adapter %q must not populate legacy Settings in new Edge-generated payload", a.GetType())
|
|
}
|
|
}
|
|
var mockFound bool
|
|
for _, a := range payload.GetAdapters() {
|
|
if a.GetType() == "mock" {
|
|
mockFound = true
|
|
if a.GetMock() == nil {
|
|
t.Fatal("mock adapter must use typed MockAdapterConfig oneof")
|
|
}
|
|
}
|
|
}
|
|
if !mockFound {
|
|
t.Fatal("expected mock adapter in payload")
|
|
}
|
|
}
|