orchestra/internal/router/router_test.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 }

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)
	}
}

// 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]string{
			"handoff_ref": handoffRef,
			"reason":      "threshold",
			"anchor_sha":  "0123456789abcdef0123456789abcdef01234567",
		})
		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.
	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")
	}
}