911 lines
29 KiB
Go
911 lines
29 KiB
Go
package service_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
toki "git.toki-labs.com/toki/proto-socket/go"
|
|
"google.golang.org/protobuf/proto"
|
|
|
|
edgenode "iop/apps/edge/internal/node"
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
func TestBuildRunRequestNormalizesSessionAndTimeout(t *testing.T) {
|
|
req, runID, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
Prompt: "hello",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
if runID == "" || req.GetRunId() != runID {
|
|
t.Fatalf("run id mismatch: runID=%q req=%q", runID, req.GetRunId())
|
|
}
|
|
if req.GetSessionId() != edgeservice.DefaultSessionID {
|
|
t.Fatalf("SessionId: got %q want %q", req.GetSessionId(), edgeservice.DefaultSessionID)
|
|
}
|
|
if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
|
|
t.Fatalf("TimeoutSec: got %d want %d", req.GetTimeoutSec(), edgeservice.DefaultTimeoutSec)
|
|
}
|
|
}
|
|
|
|
func TestBuildNodeCommandRequest_NewTypes(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
cmdType iop.NodeCommandType
|
|
prefix string
|
|
}{
|
|
{"capabilities", iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps"},
|
|
{"transport", iop.NodeCommandType_NODE_COMMAND_TYPE_TRANSPORT_STATUS, "transport"},
|
|
{"ollama api", iop.NodeCommandType_NODE_COMMAND_TYPE_OLLAMA_API, "ollama"},
|
|
}
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
req := edgeservice.BuildNodeCommandRequest(tc.cmdType, tc.prefix, "ollama", "llama3", "", 0)
|
|
if req.GetType() != tc.cmdType {
|
|
t.Errorf("Type: got %v want %v", req.GetType(), tc.cmdType)
|
|
}
|
|
if req.GetAdapter() != "ollama" {
|
|
t.Errorf("Adapter: got %q want ollama", req.GetAdapter())
|
|
}
|
|
if req.GetTarget() != "llama3" {
|
|
t.Errorf("Target: got %q want llama3", req.GetTarget())
|
|
}
|
|
if req.GetSessionId() != edgeservice.DefaultSessionID {
|
|
t.Errorf("SessionId: got %q want default", req.GetSessionId())
|
|
}
|
|
if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
|
|
t.Errorf("TimeoutSec: got %d want default", req.GetTimeoutSec())
|
|
}
|
|
if req.GetRequestId() == "" {
|
|
t.Error("RequestId: empty")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildNodeCommandRequestAndWaitTimeout(t *testing.T) {
|
|
req := edgeservice.BuildNodeCommandRequest(iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES, "caps", "ollama", "llama3", "", 0)
|
|
if req.GetSessionId() != edgeservice.DefaultSessionID {
|
|
t.Fatalf("SessionId: got %q want %q", req.GetSessionId(), edgeservice.DefaultSessionID)
|
|
}
|
|
if req.GetTimeoutSec() != edgeservice.DefaultTimeoutSec {
|
|
t.Fatalf("TimeoutSec: got %d want %d", req.GetTimeoutSec(), edgeservice.DefaultTimeoutSec)
|
|
}
|
|
if got, want := edgeservice.StatusWaitTimeout(req), 35*time.Second; got != want {
|
|
t.Fatalf("StatusWaitTimeout: got %v want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestBuildRunRequest_SessionAndBackground(t *testing.T) {
|
|
req, runID, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
SessionID: "session-a",
|
|
Background: true,
|
|
TimeoutSec: 30,
|
|
Prompt: "hello",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
if runID == "" {
|
|
t.Fatal("expected non-empty runID")
|
|
}
|
|
if req.GetSessionId() != "session-a" {
|
|
t.Errorf("SessionId: got %q want %q", req.GetSessionId(), "session-a")
|
|
}
|
|
if !req.GetBackground() {
|
|
t.Error("Background: expected true")
|
|
}
|
|
}
|
|
|
|
func TestResolveNode_RequiresExplicitSelectionForMultipleNodes(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
if _, err := svc.ResolveNode(""); err == nil {
|
|
t.Error("expected error for implicit resolve with multiple nodes")
|
|
}
|
|
if e, err := svc.ResolveNode("node-1"); err != nil || e.NodeID != "node-1" {
|
|
t.Errorf("failed explicit resolve: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestBuildRunRequestDoesNotInjectConsoleSource(t *testing.T) {
|
|
req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
Prompt: "hello",
|
|
Metadata: nil,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
if req.GetMetadata() == nil {
|
|
t.Fatal("expected non-nil metadata map")
|
|
}
|
|
if v, ok := req.GetMetadata()["source"]; ok {
|
|
t.Errorf("BuildRunRequest should not inject source, got source=%q", v)
|
|
}
|
|
}
|
|
|
|
func TestBuildRunRequestCopiesMetadata(t *testing.T) {
|
|
req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
Prompt: "hello",
|
|
Metadata: map[string]string{"request_id": "req-001", "x-custom": "value"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
if got := req.GetMetadata()["request_id"]; got != "req-001" {
|
|
t.Errorf("request_id: got %q, want %q", got, "req-001")
|
|
}
|
|
if got := req.GetMetadata()["x-custom"]; got != "value" {
|
|
t.Errorf("x-custom: got %q, want %q", got, "value")
|
|
}
|
|
}
|
|
|
|
func TestBuildRunRequestPreservesResponsesMetadataKeys(t *testing.T) {
|
|
req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama",
|
|
Prompt: "test",
|
|
Metadata: map[string]string{
|
|
"request_id": "req-001",
|
|
"task_id": "task-123",
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
for _, key := range []string{"request_id", "task_id"} {
|
|
if req.GetMetadata()[key] == "" {
|
|
t.Errorf("metadata key %q not preserved", key)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestListNodesReturnsSnapshots(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 2 {
|
|
t.Fatalf("ListNodeSnapshots: got %d want 2", len(snaps))
|
|
}
|
|
byID := map[string]edgeservice.NodeSnapshot{}
|
|
for _, s := range snaps {
|
|
byID[s.NodeID] = s
|
|
}
|
|
if got := byID["node-1"]; got.Alias != "alpha" || got.Label != "node0" {
|
|
t.Errorf("node-1 snapshot: %+v, want Label=node0", got)
|
|
}
|
|
if got := byID["node-2"]; got.Alias != "" || got.Label != "node1" {
|
|
t.Errorf("node-2 snapshot: %+v, want Label=node1", got)
|
|
}
|
|
}
|
|
|
|
func TestResolveNodeSnapshotReturnsDTO(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
snap, err := svc.ResolveNodeSnapshot("alpha")
|
|
if err != nil {
|
|
t.Fatalf("ResolveNodeSnapshot: %v", err)
|
|
}
|
|
if snap.NodeID != "node-1" || snap.Label != "node0" {
|
|
t.Errorf("snapshot: %+v", snap)
|
|
}
|
|
}
|
|
|
|
func TestBuildCancelRunRequest(t *testing.T) {
|
|
req := edgeservice.CancelRunRequest{RunID: "run-123"}
|
|
wire := edgeservice.BuildCancelRunRequest(req)
|
|
if wire.GetRunId() != req.RunID {
|
|
t.Errorf("RunId: got %q want %q", wire.GetRunId(), req.RunID)
|
|
}
|
|
}
|
|
|
|
func TestResolveNode_AllowsSingleNodeFallback(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
if e, err := svc.ResolveNode(""); err != nil || e.NodeID != "node-1" {
|
|
t.Errorf("expected single node fallback, got error: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsWithConfig(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1", Alias: "alpha"})
|
|
|
|
store := edgenode.NewNodeStore()
|
|
store.Add(&edgenode.NodeRecord{
|
|
ID: "node-1",
|
|
Alias: "alpha",
|
|
Token: "tok-1",
|
|
Runtime: config.RuntimeConf{
|
|
Concurrency: 3,
|
|
},
|
|
Adapters: config.AdaptersConf{
|
|
OllamaInstances: []config.OllamaInstanceConf{
|
|
{Name: "ollama", Enabled: true, BaseURL: "http://localhost:11434", ContextSize: 2048},
|
|
},
|
|
},
|
|
})
|
|
|
|
svc := edgeservice.New(reg, nil)
|
|
svc.SetNodeStore(store)
|
|
|
|
snaps := svc.ListNodeSnapshots()
|
|
if len(snaps) != 1 {
|
|
t.Fatalf("ListNodeSnapshots: got %d want 1", len(snaps))
|
|
}
|
|
snap := snaps[0]
|
|
if snap.NodeID != "node-1" || snap.Label != "node0" {
|
|
t.Errorf("snapshot basic info wrong: %+v", snap)
|
|
}
|
|
if snap.Config == nil {
|
|
t.Fatal("expected snap.Config to be populated")
|
|
}
|
|
if snap.Config.Runtime == nil || snap.Config.Runtime.Concurrency != 3 {
|
|
t.Errorf("unexpected runtime config: %+v", snap.Config.Runtime)
|
|
}
|
|
var ollama *iop.AdapterConfig
|
|
for _, a := range snap.Config.Adapters {
|
|
if a.Type == "ollama" {
|
|
ollama = a
|
|
break
|
|
}
|
|
}
|
|
if ollama == nil {
|
|
t.Fatal("expected ollama adapter config")
|
|
}
|
|
if !ollama.Enabled || ollama.GetOllama().GetBaseUrl() != "http://localhost:11434" {
|
|
t.Errorf("unexpected ollama config: %+v", ollama)
|
|
}
|
|
}
|
|
|
|
func TestListNodeSnapshotsIncludesLifecycleState(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-generic"})
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-explicit",
|
|
LifecycleState: "registering",
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
byID := map[string]edgeservice.NodeSnapshot{}
|
|
for _, s := range svc.ListNodeSnapshots() {
|
|
byID[s.NodeID] = s
|
|
}
|
|
|
|
generic, ok := byID["node-generic"]
|
|
if !ok {
|
|
t.Fatal("node-generic snapshot missing")
|
|
}
|
|
if generic.LifecycleState != edgenode.LifecycleConnected {
|
|
t.Errorf("generic LifecycleState: got %q want %q", generic.LifecycleState, edgenode.LifecycleConnected)
|
|
}
|
|
|
|
explicit, ok := byID["node-explicit"]
|
|
if !ok {
|
|
t.Fatal("node-explicit snapshot missing")
|
|
}
|
|
if explicit.LifecycleState != "registering" {
|
|
t.Errorf("explicit LifecycleState: got %q want %q", explicit.LifecycleState, "registering")
|
|
}
|
|
}
|
|
|
|
func TestResolveNodeSnapshotCarriesLifecycle(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-explicit",
|
|
Alias: "explicit-1",
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
snap, err := svc.ResolveNodeSnapshot("explicit-1")
|
|
if err != nil {
|
|
t.Fatalf("ResolveNodeSnapshot: %v", err)
|
|
}
|
|
if snap.LifecycleState != edgenode.LifecycleConnected {
|
|
t.Errorf("LifecycleState: got %q want %q", snap.LifecycleState, edgenode.LifecycleConnected)
|
|
}
|
|
}
|
|
|
|
func TestServiceCapabilitiesAreNodeScoped(t *testing.T) {
|
|
for _, state := range []string{"", "connected", "deploying", "running", "busy", "error", "failed", "unknown-state"} {
|
|
t.Run(fmt.Sprintf("state-%s", state), func(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-1",
|
|
Alias: "gen-node",
|
|
})
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-2",
|
|
Alias: "node2",
|
|
LifecycleState: state,
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
caps := svc.GetCapabilities()
|
|
if len(caps) != 1 {
|
|
t.Fatalf("expected 1 capability, got %d", len(caps))
|
|
}
|
|
capMap := map[string]*iop.EdgeCapabilitySummary{}
|
|
for _, c := range caps {
|
|
capMap[c.Kind] = c
|
|
}
|
|
if c := capMap["run"]; c == nil || !c.Available || c.Status != "ready" {
|
|
t.Errorf("invalid run capability: %+v", c)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestServiceExecuteCommand(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-1",
|
|
Alias: "deployer",
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
t.Run("health.check", func(t *testing.T) {
|
|
var events []*iop.EdgeCommandEvent
|
|
onEvent := func(e *iop.EdgeCommandEvent) {
|
|
events = append(events, e)
|
|
}
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-hc",
|
|
Operation: "health.check",
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand error: %v", err)
|
|
}
|
|
if resp.Status != "completed" {
|
|
t.Errorf("status: got %q want completed", resp.Status)
|
|
}
|
|
if len(events) != 2 {
|
|
t.Fatalf("expected 2 events, got %d", len(events))
|
|
}
|
|
if events[0].Phase != "started" || events[1].Phase != "completed" {
|
|
t.Errorf("unexpected event sequence: %+v", events)
|
|
}
|
|
})
|
|
|
|
t.Run("node.status found", func(t *testing.T) {
|
|
var events []*iop.EdgeCommandEvent
|
|
onEvent := func(e *iop.EdgeCommandEvent) {
|
|
events = append(events, e)
|
|
}
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-status",
|
|
Operation: "node.status",
|
|
TargetSelector: "deployer",
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand error: %v", err)
|
|
}
|
|
if resp.Status != "completed" {
|
|
t.Errorf("status: got %q want completed", resp.Status)
|
|
}
|
|
if len(events) != 2 {
|
|
t.Fatalf("expected 2 events, got %d", len(events))
|
|
}
|
|
if events[0].Phase != "started" || events[1].Phase != "completed" {
|
|
t.Errorf("unexpected event sequence: %+v", events)
|
|
}
|
|
})
|
|
|
|
t.Run("node.status missing", func(t *testing.T) {
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-status",
|
|
Operation: "node.status",
|
|
TargetSelector: "missing-agent",
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.Status != "error" || resp.Error == "" {
|
|
t.Errorf("expected error status, got: %+v", resp)
|
|
}
|
|
})
|
|
|
|
t.Run("provider.command does not fake success on nil client", func(t *testing.T) {
|
|
var events []*iop.EdgeCommandEvent
|
|
onEvent := func(e *iop.EdgeCommandEvent) {
|
|
events = append(events, e)
|
|
}
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-run",
|
|
Operation: "provider.command",
|
|
TargetSelector: "deployer",
|
|
Parameters: map[string]string{
|
|
"command": "capabilities",
|
|
},
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.Status != "error" {
|
|
t.Errorf("expected error status, got: %q", resp.Status)
|
|
}
|
|
if len(events) < 2 {
|
|
t.Fatalf("expected at least started and failed events, got %d", len(events))
|
|
}
|
|
if events[0].Phase != "started" || events[len(events)-1].Phase != "failed" {
|
|
t.Errorf("unexpected event sequence: %+v", events)
|
|
}
|
|
})
|
|
|
|
t.Run("provider.command missing command param", func(t *testing.T) {
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-run",
|
|
Operation: "provider.command",
|
|
TargetSelector: "deployer",
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.Status != "unsupported" {
|
|
t.Errorf("expected unsupported status, got %q", resp.Status)
|
|
}
|
|
})
|
|
|
|
t.Run("unsupported operation", func(t *testing.T) {
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-bad",
|
|
Operation: "bad.op",
|
|
}
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, func(*iop.EdgeCommandEvent) {})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.Status != "unsupported" {
|
|
t.Errorf("status: got %q want unsupported", resp.Status)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestServiceExecuteProviderCommandRequiresExplicitSelector(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "only-ready-node"})
|
|
svc := edgeservice.New(reg, nil)
|
|
var events []*iop.EdgeCommandEvent
|
|
|
|
response, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
|
|
CommandId: "missing-selector",
|
|
Operation: "provider.command",
|
|
Parameters: map[string]string{"command": "capabilities"},
|
|
}, func(event *iop.EdgeCommandEvent) { events = append(events, event) })
|
|
if err == nil || !strings.Contains(err.Error(), "target_selector is required") {
|
|
t.Fatalf("expected explicit selector error, got response=%+v err=%v", response, err)
|
|
}
|
|
if response != nil {
|
|
t.Fatalf("expected no response before command start, got %+v", response)
|
|
}
|
|
if len(events) != 0 {
|
|
t.Fatalf("expected no events before command start, got %+v", events)
|
|
}
|
|
}
|
|
|
|
func TestServiceExecuteProviderCommandRejectsUnsupportedBeforeStart(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "only-ready-node"})
|
|
svc := edgeservice.New(reg, nil)
|
|
var events []*iop.EdgeCommandEvent
|
|
|
|
response, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
|
|
CommandId: "unsupported-command",
|
|
Operation: "provider.command",
|
|
TargetSelector: "only-ready-node",
|
|
Parameters: map[string]string{"command": "shell.exec"},
|
|
}, func(event *iop.EdgeCommandEvent) { events = append(events, event) })
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand: %v", err)
|
|
}
|
|
if response.GetStatus() != "unsupported" {
|
|
t.Fatalf("expected unsupported response, got %+v", response)
|
|
}
|
|
if len(events) != 0 {
|
|
t.Fatalf("unsupported command emitted events before start: %+v", events)
|
|
}
|
|
}
|
|
|
|
func TestServiceExecuteProviderCommandTerminalEvents(t *testing.T) {
|
|
for _, command := range []string{"capabilities", "transport_status", "ollama_api"} {
|
|
for _, nodeError := range []bool{false, true} {
|
|
t.Run(fmt.Sprintf("%s/error=%t", command, nodeError), func(t *testing.T) {
|
|
svc := newProviderCommandTestService(t, nodeError)
|
|
params := map[string]string{"command": command, "adapter": "ollama", "target": "llama3"}
|
|
if command == "ollama_api" {
|
|
params["method"] = "GET"
|
|
params["path"] = "/api/tags"
|
|
}
|
|
var events []*iop.EdgeCommandEvent
|
|
response, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
|
|
CommandId: "terminal-events",
|
|
Operation: "provider.command",
|
|
TargetSelector: "provider-node",
|
|
Parameters: params,
|
|
}, func(event *iop.EdgeCommandEvent) { events = append(events, event) })
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand: %v", err)
|
|
}
|
|
if len(events) != 2 || events[0].GetPhase() != "started" {
|
|
t.Fatalf("expected exactly started plus terminal event, got %+v", events)
|
|
}
|
|
wantTerminal := "completed"
|
|
if nodeError {
|
|
wantTerminal = "failed"
|
|
if response.GetStatus() != "error" {
|
|
t.Fatalf("error response: got %+v", response)
|
|
}
|
|
} else if response.GetStatus() != "completed" {
|
|
t.Fatalf("completed response: got %+v", response)
|
|
}
|
|
if events[1].GetPhase() != wantTerminal {
|
|
t.Fatalf("terminal phase: got %q want %q", events[1].GetPhase(), wantTerminal)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
func newProviderCommandTestService(t *testing.T, nodeError bool) *edgeservice.Service {
|
|
t.Helper()
|
|
edgeConn, nodeConn := net.Pipe()
|
|
t.Cleanup(func() {
|
|
_ = edgeConn.Close()
|
|
_ = nodeConn.Close()
|
|
})
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
|
|
message := &iop.NodeCommandRequest{}
|
|
return message, proto.Unmarshal(b, message)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
|
|
message := &iop.NodeCommandResponse{}
|
|
return message, proto.Unmarshal(b, message)
|
|
},
|
|
}
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
toki.AddRequestListenerTyped(&nodeClient.Communicator, func(request *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
|
|
response := &iop.NodeCommandResponse{RequestId: request.GetRequestId(), Type: request.GetType()}
|
|
if nodeError {
|
|
response.Error = "provider command failed"
|
|
return response, nil
|
|
}
|
|
response.Result = map[string]string{"status_code": "200", "result": "ok"}
|
|
return response, nil
|
|
})
|
|
registry := edgenode.NewRegistry()
|
|
registry.Register(&edgenode.NodeEntry{NodeID: "provider-node", Client: edgeClient})
|
|
return edgeservice.New(registry, nil)
|
|
}
|
|
|
|
func TestServiceExecuteCommandHealthCheckDerivesFromStatus(t *testing.T) {
|
|
t.Run("with nodes summarizes registry and capabilities", func(t *testing.T) {
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-1"})
|
|
reg.Register(&edgenode.NodeEntry{NodeID: "node-2"})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
// HealthStatus is the Edge-owned boundary the command must read from.
|
|
status := svc.HealthStatus()
|
|
if len(status.Nodes) != 2 {
|
|
t.Fatalf("HealthStatus nodes: got %d want 2", len(status.Nodes))
|
|
}
|
|
if len(status.Capabilities) != 1 {
|
|
t.Fatalf("HealthStatus capabilities: got %d want 1", len(status.Capabilities))
|
|
}
|
|
|
|
var events []*iop.EdgeCommandEvent
|
|
resp, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-hc",
|
|
Operation: "health.check",
|
|
}, func(e *iop.EdgeCommandEvent) { events = append(events, e) })
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand error: %v", err)
|
|
}
|
|
if resp.Status != "completed" {
|
|
t.Errorf("status: got %q want completed", resp.Status)
|
|
}
|
|
if !strings.Contains(resp.Summary, "2 node(s)") || !strings.Contains(resp.Summary, "1 capability(ies)") {
|
|
t.Errorf("summary not derived from status: got %q", resp.Summary)
|
|
}
|
|
if len(events) != 2 || events[1].Phase != "completed" {
|
|
t.Fatalf("unexpected events: %+v", events)
|
|
}
|
|
if events[1].Summary != resp.Summary {
|
|
t.Errorf("completed event summary %q != response summary %q", events[1].Summary, resp.Summary)
|
|
}
|
|
})
|
|
|
|
t.Run("without nodes reports empty status", func(t *testing.T) {
|
|
svc := edgeservice.New(edgenode.NewRegistry(), nil)
|
|
resp, err := svc.ExecuteCommand(context.Background(), &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-hc",
|
|
Operation: "health.check",
|
|
}, func(*iop.EdgeCommandEvent) {})
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand error: %v", err)
|
|
}
|
|
if !strings.Contains(resp.Summary, "0 node(s)") || !strings.Contains(resp.Summary, "0 capability(ies)") {
|
|
t.Errorf("empty status summary wrong: got %q", resp.Summary)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestServiceExecuteCommandProviderCommandCapabilitiesSuccess(t *testing.T) {
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandResponse{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
|
|
if req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
|
|
return nil, fmt.Errorf("unexpected command type: %v", req.GetType())
|
|
}
|
|
if req.GetAdapter() != "ollama" || req.GetTarget() != "llama3" || req.GetSessionId() != "sess-success" {
|
|
return nil, fmt.Errorf("unexpected adapter/target/session: %s/%s/%s", req.GetAdapter(), req.GetTarget(), req.GetSessionId())
|
|
}
|
|
return &iop.NodeCommandResponse{
|
|
RequestId: req.RequestId,
|
|
Type: req.Type,
|
|
Result: map[string]string{"caps": "ollama,llama3"},
|
|
}, nil
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-success-1",
|
|
Alias: "deployer",
|
|
Client: edgeClient,
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
var events []*iop.EdgeCommandEvent
|
|
onEvent := func(e *iop.EdgeCommandEvent) {
|
|
events = append(events, e)
|
|
}
|
|
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-success-1",
|
|
Operation: "provider.command",
|
|
TargetSelector: "deployer",
|
|
Parameters: map[string]string{
|
|
"command": "capabilities",
|
|
"adapter": "ollama",
|
|
"target": "llama3",
|
|
"session_id": "sess-success",
|
|
},
|
|
}
|
|
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand: %v", err)
|
|
}
|
|
|
|
if resp.Status != "completed" {
|
|
t.Errorf("status: got %q want completed", resp.Status)
|
|
}
|
|
if !strings.Contains(resp.Summary, "caps=ollama,llama3") {
|
|
t.Errorf("summary: got %q, expected containing caps=ollama,llama3", resp.Summary)
|
|
}
|
|
|
|
if len(events) < 2 {
|
|
t.Fatalf("expected at least 2 events, got %d", len(events))
|
|
}
|
|
if events[0].Phase != "started" || events[len(events)-1].Phase != "completed" {
|
|
t.Errorf("unexpected event sequence: %+v", events)
|
|
}
|
|
if events[len(events)-1].Summary != "Command capabilities completed successfully" {
|
|
t.Errorf("unexpected event summary: got %q, expected 'Command capabilities completed successfully'", events[len(events)-1].Summary)
|
|
}
|
|
}
|
|
|
|
func TestServiceExecuteCommandProviderCommandCapabilitiesNodeError(t *testing.T) {
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandResponse{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
|
|
if req.GetSessionId() != "sess-error" {
|
|
return nil, fmt.Errorf("unexpected session_id: %s", req.GetSessionId())
|
|
}
|
|
return &iop.NodeCommandResponse{
|
|
RequestId: req.RequestId,
|
|
Type: req.Type,
|
|
Error: "node failure",
|
|
}, nil
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-error-1",
|
|
Alias: "deployer",
|
|
Client: edgeClient,
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
var events []*iop.EdgeCommandEvent
|
|
onEvent := func(e *iop.EdgeCommandEvent) {
|
|
events = append(events, e)
|
|
}
|
|
|
|
req := &iop.EdgeCommandRequest{
|
|
CommandId: "cmd-error-1",
|
|
Operation: "provider.command",
|
|
TargetSelector: "deployer",
|
|
Parameters: map[string]string{
|
|
"command": "capabilities",
|
|
"adapter": "ollama",
|
|
"target": "llama3",
|
|
"session_id": "sess-error",
|
|
},
|
|
}
|
|
|
|
resp, err := svc.ExecuteCommand(context.Background(), req, onEvent)
|
|
if err != nil {
|
|
t.Fatalf("ExecuteCommand: %v", err)
|
|
}
|
|
|
|
if resp.Status != "error" {
|
|
t.Errorf("status: got %q want error", resp.Status)
|
|
}
|
|
if !strings.Contains(resp.Error, "node reported error: node failure") {
|
|
t.Errorf("error: got %q, expected containing node failure", resp.Error)
|
|
}
|
|
|
|
if len(events) < 2 {
|
|
t.Fatalf("expected at least 2 events, got %d", len(events))
|
|
}
|
|
if events[0].Phase != "started" || events[len(events)-1].Phase != "failed" {
|
|
t.Errorf("unexpected event sequence: %+v", events)
|
|
}
|
|
expectedSummary := "Command capabilities failed: node reported error: node failure"
|
|
if events[len(events)-1].Summary != expectedSummary {
|
|
t.Errorf("unexpected event summary: got %q, expected %q", events[len(events)-1].Summary, expectedSummary)
|
|
}
|
|
}
|
|
|
|
func TestServiceCapabilitiesPreservesProviderSnapshots(t *testing.T) {
|
|
edgeConn, nodeConn := net.Pipe()
|
|
defer edgeConn.Close()
|
|
defer nodeConn.Close()
|
|
|
|
parserMap := toki.ParserMap{
|
|
toki.TypeNameOf(&iop.NodeCommandRequest{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandRequest{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
toki.TypeNameOf(&iop.NodeCommandResponse{}): func(b []byte) (proto.Message, error) {
|
|
m := &iop.NodeCommandResponse{}
|
|
return m, proto.Unmarshal(b, m)
|
|
},
|
|
}
|
|
|
|
edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
|
|
nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
|
|
|
|
toki.AddRequestListenerTyped(&nodeClient.Communicator, func(req *iop.NodeCommandRequest) (*iop.NodeCommandResponse, error) {
|
|
if req.GetType() != iop.NodeCommandType_NODE_COMMAND_TYPE_CAPABILITIES {
|
|
return nil, fmt.Errorf("unexpected command type: %v", req.GetType())
|
|
}
|
|
return &iop.NodeCommandResponse{
|
|
RequestId: req.RequestId,
|
|
Type: req.Type,
|
|
Result: map[string]string{"caps": "ollama,llama3"},
|
|
ProviderSnapshots: []*iop.ProviderSnapshot{
|
|
{
|
|
Adapter: "ollama",
|
|
Status: "available",
|
|
Capacity: 4,
|
|
InFlight: 2,
|
|
Queued: 3,
|
|
},
|
|
},
|
|
}, nil
|
|
})
|
|
|
|
reg := edgenode.NewRegistry()
|
|
reg.Register(&edgenode.NodeEntry{
|
|
NodeID: "node-caps-1",
|
|
Alias: "deployer",
|
|
Client: edgeClient,
|
|
})
|
|
svc := edgeservice.New(reg, nil)
|
|
|
|
view, err := svc.Capabilities(context.Background(), edgeservice.NodeCommandRequestSpec{
|
|
NodeRef: "deployer",
|
|
Adapter: "ollama",
|
|
Target: "llama3",
|
|
SessionID: "sess-caps",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Capabilities: %v", err)
|
|
}
|
|
|
|
if len(view.ProviderSnapshots) != 1 {
|
|
t.Fatalf("expected 1 provider snapshot, got %d", len(view.ProviderSnapshots))
|
|
}
|
|
snap := view.ProviderSnapshots[0]
|
|
if snap.Adapter != "ollama" || snap.Status != "available" || snap.Capacity != 4 || snap.InFlight != 2 || snap.Queued != 3 {
|
|
t.Fatalf("unexpected snapshot contents: %+v", snap)
|
|
}
|
|
}
|
|
|
|
func TestBuildRunRequestPreservesEstimatedTokensAndContextClass(t *testing.T) {
|
|
req, _, err := edgeservice.BuildRunRequest(edgeservice.SubmitRunRequest{
|
|
Adapter: "ollama",
|
|
Target: "llama",
|
|
Prompt: "test",
|
|
EstimatedInputTokens: 150000,
|
|
ContextClass: "long",
|
|
Metadata: map[string]string{"estimated_input_tokens": "150000", "context_class": "long"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("BuildRunRequest: %v", err)
|
|
}
|
|
// BuildRunRequest does not carry EstimatedInputTokens/ContextClass in the protobuf,
|
|
// so we verify that the metadata copy is preserved.
|
|
if req.GetMetadata()["estimated_input_tokens"] != "150000" {
|
|
t.Fatalf("metadata estimated_input_tokens: got %q", req.GetMetadata()["estimated_input_tokens"])
|
|
}
|
|
if req.GetMetadata()["context_class"] != "long" {
|
|
t.Fatalf("metadata context_class: got %q", req.GetMetadata()["context_class"])
|
|
}
|
|
}
|