kami
208 lines
8.2 KiB
Go
208 lines
8.2 KiB
Go
package router
|
|
|
|
import (
|
|
"encoding/json"
|
|
"orchestra/internal/authz"
|
|
"orchestra/internal/domain"
|
|
"orchestra/internal/registry"
|
|
"orchestra/internal/store"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
type reachable struct{}
|
|
|
|
func (reachable) Reachable(string, time.Duration) bool { return true }
|
|
|
|
type projectAvailability struct{ projects map[string]bool }
|
|
|
|
func (p projectAvailability) Available(registry.Herdr) bool { return true }
|
|
func (p projectAvailability) Supports(h registry.Herdr, project string) bool {
|
|
return p.projects[h.ID+"/"+project]
|
|
}
|
|
|
|
func TestAssignsByAffinityCapabilityAndConcurrency(t *testing.T) {
|
|
s, err := store.Open(t.TempDir())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
r, err := registry.New(registry.Config{
|
|
Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
|
|
Machines: []registry.Machine{{ID: "m", Address: "unused"}},
|
|
Herdrs: []registry.Herdr{{ID: "h", MachineID: "m", Capabilities: []string{"go"}, Concurrency: 1}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
makeTask := func(id string) {
|
|
b, _ := json.Marshal(map[string]any{"source": "test", "external_id": id, "project": "p", "capability": []string{"go"}})
|
|
if err := s.Append(domain.Event{ID: id, TaskID: id, Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
makeTask("a")
|
|
makeTask("b")
|
|
rt := Router{Store: s, Registry: r, Reachability: reachable{}}
|
|
got, err := rt.AssignPending()
|
|
if err != nil || len(got) != 1 {
|
|
t.Fatalf("assigned %d events, err=%v", len(got), err)
|
|
}
|
|
// Store.Tasks() ranges a map, so its order is randomized per call. This
|
|
// assertion used to index two *separate* Tasks() calls, and failed
|
|
// whenever the two orderings disagreed — the flake was in the test, not
|
|
// in assignment. Snapshot once, and assert the actual invariant:
|
|
// concurrency 1 means exactly one of the two tasks is leased.
|
|
leased := 0
|
|
for _, tk := range s.Tasks() {
|
|
if tk.State == domain.StateLeased {
|
|
leased++
|
|
}
|
|
}
|
|
if leased != 1 {
|
|
t.Fatalf("leased %d tasks, want exactly 1 (concurrency 1); got=%+v", leased, got)
|
|
}
|
|
}
|
|
|
|
func TestAssignPendingRequiresWorkerProjectSupport(t *testing.T) {
|
|
s, err := store.Open(t.TempDir())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
r, err := registry.New(registry.Config{Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}}, Machines: []registry.Machine{{ID: "m", Address: "unused"}}, Herdrs: []registry.Herdr{{ID: "h", MachineID: "m", Concurrency: 1}}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "project-check", "project": "p"})
|
|
if err := s.Append(domain.Event{ID: "create", TaskID: "t", Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rt := Router{Store: s, Registry: r, Reachability: reachable{}, Availability: projectAvailability{projects: map[string]bool{}}}
|
|
if got, err := rt.AssignPending(); err != nil || len(got) != 0 {
|
|
t.Fatalf("unsupported project lease = %#v, %v", got, err)
|
|
}
|
|
if got, _ := s.Task("t"); got.State != domain.StateQueued {
|
|
t.Fatalf("unsupported project state=%s", got.State)
|
|
}
|
|
rt.Availability = projectAvailability{projects: map[string]bool{"h/p": true}}
|
|
if got, err := rt.AssignPending(); err != nil || len(got) != 1 {
|
|
t.Fatalf("supported project lease = %#v, %v", got, err)
|
|
}
|
|
}
|
|
|
|
// TestRotationDoesNotCountAgainstRetryLimit guards B4: rotation is
|
|
// TaskReleased carrying a valid handoff_ref (spec §5.3: "rotation =
|
|
// intra-task lease transfer"), never a failure. A task healthy enough to
|
|
// rotate repeatedly must survive past MaxAttempts, which is a retry policy
|
|
// for genuine failures (expiry/crash releases without a handoff_ref), not
|
|
// for lease transfers.
|
|
func TestRotationDoesNotCountAgainstRetryLimit(t *testing.T) {
|
|
s, err := store.Open(t.TempDir())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
r, err := registry.New(registry.Config{
|
|
Projects: []registry.Project{{ID: "p", MachineAffinity: []string{"m"}}},
|
|
Machines: []registry.Machine{{ID: "m", Address: "unused"}},
|
|
Herdrs: []registry.Herdr{{ID: "h", MachineID: "m", Capabilities: []string{"go"}, Concurrency: 1}},
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
b, _ := json.Marshal(map[string]any{"source": "test", "external_id": "a", "project": "p", "capability": []string{"go"}})
|
|
if err := s.Append(domain.Event{ID: "a", TaskID: "a", Type: "TaskCreated", Version: 1, Payload: b, Surface: string(authz.System)}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
rt := Router{Store: s, Registry: r, Reachability: reachable{}, Retry: RetryPolicy{MaxAttempts: 3}}
|
|
handoffRef, err := s.PutArtifact([]byte("handoff"))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
for i := 0; i < 5; i++ {
|
|
got, err := rt.AssignPending()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
task, ok := s.Task("a")
|
|
if !ok {
|
|
t.Fatal("task missing")
|
|
}
|
|
if task.State == domain.StateFailed {
|
|
t.Fatalf("task failed after %d rotations, retry limit wrongly counted rotation as a failure", i)
|
|
}
|
|
if task.State != domain.StateLeased {
|
|
if len(got) == 0 {
|
|
t.Fatalf("round %d: task not leased and nothing assigned (state=%v)", i, task.State)
|
|
}
|
|
continue
|
|
}
|
|
rb, _ := json.Marshal(map[string]any{
|
|
"handoff_ref": handoffRef,
|
|
"reason": "threshold",
|
|
"anchor_sha": "0123456789abcdef0123456789abcdef01234567",
|
|
"harness_id": task.Lease.HarnessID,
|
|
"lease_epoch": task.Lease.Epoch,
|
|
"expected_version": task.Version,
|
|
})
|
|
release := domain.Event{ID: domain.NewID(), Type: "TaskReleased", TaskID: "a", Version: task.Version + 1, Payload: rb, Surface: string(authz.System)}
|
|
if err := s.Append(release); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := rt.HandleEvent(release); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
task, _ := s.Task("a")
|
|
if task.State == domain.StateFailed {
|
|
t.Fatal("task failed after 5 rotations, want still alive")
|
|
}
|
|
}
|
|
|
|
// TestQuotaWindowsAreIndependent proves the 5-hour rolling window and the
|
|
// weekly window (spec §7.2, §9 item 1) are each conservative-80%-full gates
|
|
// on their own — a harness can be fine on one window and excluded by the
|
|
// other, and receipts outside a window must not count toward it.
|
|
func TestQuotaWindowsAreIndependent(t *testing.T) {
|
|
s, err := store.Open(t.TempDir())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
now := time.Now().UTC()
|
|
report := func(at time.Time, consumed float64) {
|
|
p, _ := json.Marshal(map[string]any{"harness_id": "h1", "consumed": consumed})
|
|
if err := s.Append(domain.Event{ID: domain.NewID(), Type: "QuotaReported", TaskID: "quota", Version: 1, Surface: string(authz.System), Payload: p, At: at}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
// Case 1: only weekly limit configured. A receipt older than 5h but
|
|
// within the week still counts toward the weekly gate.
|
|
report(now.Add(-6*time.Hour), 85)
|
|
weeklyOnly := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {Weekly: 100}}, Now: func() time.Time { return now }}
|
|
if weeklyOnly.Available(registry.Herdr{ID: "h1"}) {
|
|
t.Fatal("weekly window should be exhausted at 85/100 (>=80%)")
|
|
}
|
|
|
|
// Case 2: only a 5h limit configured. The same 6h-old receipt is outside
|
|
// the 5h window and must not count; a fresh zero receipt proves the native
|
|
// usage source is known for this window.
|
|
report(now, 0)
|
|
fiveHourOnly := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {FiveHour: 100}}, Now: func() time.Time { return now }}
|
|
if !fiveHourOnly.Available(registry.Herdr{ID: "h1"}) {
|
|
t.Fatal("receipt outside the 5h window incorrectly counted against it")
|
|
}
|
|
|
|
// Case 3: a fresh receipt inside the 5h window trips the 5h gate even
|
|
// though the weekly gate (fed by both receipts) also trips — both are
|
|
// independently enforced, and either failing excludes the harness.
|
|
report(now.Add(-time.Minute), 90)
|
|
both := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"h1": {FiveHour: 100, Weekly: 500}}, Now: func() time.Time { return now }}
|
|
if both.Available(registry.Herdr{ID: "h1"}) {
|
|
t.Fatal("5h window should be exhausted at 90/100 (>=80%) regardless of weekly headroom")
|
|
}
|
|
unknown := QuotaAvailability{Store: s, Limits: map[string]QuotaWindowLimits{"unknown": {FiveHour: 100}}, Now: func() time.Time { return now }}
|
|
if unknown.Available(registry.Herdr{ID: "unknown"}) {
|
|
t.Fatal("bounded harness without a native usage receipt must fail closed")
|
|
}
|
|
}
|