61 lines
2.5 KiB
Go
61 lines
2.5 KiB
Go
package agent_execution_test
|
|
|
|
import (
|
|
"context"
|
|
"strconv"
|
|
"testing"
|
|
|
|
"git.cin.su/m/Assistant/contracts/work"
|
|
"git.cin.su/m/Assistant/modules/work/application"
|
|
)
|
|
|
|
type graphStore struct {
|
|
graphs map[string]work.WorkGraph
|
|
queued []string
|
|
intents []string
|
|
}
|
|
|
|
func (s *graphStore) SaveWorkGraph(_ context.Context, graph work.WorkGraph) error {
|
|
if s.graphs == nil {
|
|
s.graphs = map[string]work.WorkGraph{}
|
|
}
|
|
s.graphs[graph.ID+":"+itoa(graph.Revision)] = graph
|
|
return nil
|
|
}
|
|
func (s *graphStore) LoadWorkGraph(_ context.Context, id string, revision int64) (work.WorkGraph, error) {
|
|
return s.graphs[id+":"+itoa(revision)], nil
|
|
}
|
|
func (s *graphStore) SaveDispatchIntent(_ context.Context, intent work.AgentDispatchIntent) error {
|
|
s.intents = append(s.intents, intent.DispatchID)
|
|
return nil
|
|
}
|
|
func (s *graphStore) EnqueueDispatch(_ context.Context, item work.AgentDispatchQueueItem) error {
|
|
s.queued = append(s.queued, item.DispatchID)
|
|
return nil
|
|
}
|
|
|
|
func itoa(value int64) string {
|
|
return strconv.FormatInt(value, 10)
|
|
}
|
|
|
|
func TestGraphRuntimeSchedulesFanOutAndCompletesFanIn(t *testing.T) {
|
|
store := &graphStore{}
|
|
graph, ready, err := application.AdmitAndSchedule(context.Background(), store, application.GraphScheduleRequest{Graph: work.WorkGraph{ID: "g", Revision: 1, Budget: work.GraphBudget{MaxParallel: 2, MaxAttempts: 1}, Nodes: []work.GraphNode{{ID: "a", DispatchID: "d-a", IdempotencyKey: "a", State: work.GraphNodeReady}, {ID: "b", DispatchID: "d-b", IdempotencyKey: "b", State: work.GraphNodeReady}, {ID: "join", DispatchID: "d-join", IdempotencyKey: "join", State: work.GraphNodeBlocked}}, Edges: []work.GraphEdge{{From: "a", To: "join"}, {From: "b", To: "join"}}}, Intents: []work.AgentDispatchIntent{{DispatchID: "d-a"}, {DispatchID: "d-b"}, {DispatchID: "d-join"}}, Now: 1})
|
|
if err != nil || len(ready) != 2 || len(store.queued) != 2 {
|
|
t.Fatalf("graph=%+v ready=%v queued=%v err=%v", graph, ready, store.queued, err)
|
|
}
|
|
if _, err := application.CompleteGraphNode(context.Background(), store, "g", 2, "a", true); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := application.CompleteGraphNode(context.Background(), store, "g", 3, "b", true); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(store.queued) != 3 || store.queued[2] != "d-join" {
|
|
t.Fatalf("fan-in dispatch was not enqueued: %v", store.queued)
|
|
}
|
|
completed, err := application.CompleteGraphNode(context.Background(), store, "g", 4, "join", true)
|
|
if err != nil || completed.State != work.GraphCompleted {
|
|
t.Fatalf("completed=%+v err=%v", completed, err)
|
|
}
|
|
}
|