iop/apps/edge/internal/service/model_queue_admission_test.go

package service

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"testing"
	"time"

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// TestModelQueueFIFOOrdering verifies that queued items are dispatched in FIFO
// order when a slot becomes available.
func TestModelQueueFIFOOrdering(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-q1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-q1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(store)

	// Fill the only capacity slot.
	first, err := m.admit(context.Background(), "g-fifo", "", "", cands, defPolicy, nil, false, false)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Manually insert two items in known order so we can verify FIFO.
	item1 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	item2 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-fifo", groupPolicy{})
	g.queue = append(g.queue, item1, item2)
	m.mu.Unlock()

	// Release one slot — item1 (head) must be dispatched first.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "", false)
	m.mu.Unlock()

	select {
	case res := <-item1.waitCh:
		if res.err != nil {
			t.Fatalf("item1 expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.entry.NodeID != "node-q1" {
			t.Fatalf("item1: unexpected candidate %v", res.candidate)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item1 was not dispatched after slot release")
	}

	// item2 must still be waiting (no slot available yet).
	select {
	case <-item2.waitCh:
		t.Fatal("item2 should not have been dispatched yet")
	default:
	}

	// Release the slot item1 holds so item2 gets dispatched.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1", "", false)
	m.mu.Unlock()

	select {
	case res := <-item2.waitCh:
		if res.err != nil {
			t.Fatalf("item2 expected dispatch, got error: %v", res.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item2 was not dispatched after second slot release")
	}
}

// TestModelQueueOverflow verifies that admit rejects requests when the queue
// is at max_queue capacity.
func TestModelQueueOverflow(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ov1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-ov1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill the node's capacity and set max_queue=1.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-overflow", groupPolicy{})
	g.policy.maxQueue = 1
	g.inflight["node-ov1"] = 1
	m.mu.Unlock()

	// Queue one item — should succeed.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	m.groups["g-overflow"].queue = append(m.groups["g-overflow"].queue, item)
	m.mu.Unlock()

	// Second admit should fail immediately with errQueueFull.
	_, err := m.admit(context.Background(), "g-overflow", "", "", cands, groupPolicy{}, nil, false, false)
	if !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull, got: %v", err)
	}
}

// TestModelQueueTimeout verifies that queued items receive a timeout error
// when no slot becomes available before the deadline.
func TestModelQueueTimeout(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-to1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-to1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill capacity and set a very short queue timeout.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-timeout", groupPolicy{})
	g.policy.queueTimeout = 20 * time.Millisecond
	g.inflight["node-to1"] = 1
	m.mu.Unlock()

	start := time.Now()
	_, err := m.admit(context.Background(), "g-timeout", "", "", cands, groupPolicy{}, nil, false, false)
	elapsed := time.Since(start)

	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected errQueueTimeout, got: %v", err)
	}
	if elapsed < 15*time.Millisecond {
		t.Fatalf("timed out too fast: %v", elapsed)
	}
}

// TestModelQueueTerminalReleaseDispatchesNext verifies that a terminal run event
// releases the in-flight slot and dispatches the next queued item.
func TestModelQueueTerminalReleaseDispatchesNext(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-tr1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-tr1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(store)

	// Fill capacity and record inflight.
	selected, err := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy, nil, false, false)
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	m.trackLease(selected.leaseID, "run-tr-001")

	// Queue a second item in a goroutine.
	resultCh := make(chan admitResult, 1)
	go func() {
		n, e := m.admit(context.Background(), "g-tr", "", "", cands, defPolicy, nil, false, false)
		resultCh <- admitResult{candidate: n, err: e}
	}()

	waitForQueueLen(t, m, "g-tr", 1)

	// Terminal event for the first run — should unblock the queued admit.
	for _, termType := range []string{"complete", "error", "cancelled"} {
		t.Run(termType, func(t *testing.T) {
			// Reset state for each sub-test.
			m.mu.Lock()
			g := m.getOrCreateGroupLocked("g-tr-"+termType, groupPolicy{})
			g.inflight["node-tr1"] = 1
			item := &queueItem{
				candidates: cands,
				waitCh:     make(chan admitResult, 1),
				deadline:   time.Now().Add(2 * time.Second),
			}
			g.queue = []*queueItem{item}
			runID := "run-tr-" + termType
			m.mu.Unlock()
			trackTestLease(m, "g-tr-"+termType, runID, "node-tr1", "", false)

			// The transport lifecycle hook calls this directly; the event bus is
			// not in the release path.
			m.releaseRun(runID, termType)

			select {
			case res := <-item.waitCh:
				if res.err != nil {
					t.Fatalf("expected dispatch, got error: %v", res.err)
				}
				if res.candidate == nil {
					t.Fatal("expected non-nil candidate")
				}
			case <-time.After(500 * time.Millisecond):
				t.Fatalf("timeout: queued item not dispatched after %q terminal event", termType)
			}
		})
	}

	// Cleanup the goroutine from the outer admit.
	m.mu.Lock()
	if g, ok := m.groups["g-tr"]; ok && len(g.queue) > 0 {
		g.inflight["node-tr1"] = 1
		m.pumpAllLocked()
	}
	m.mu.Unlock()
	select {
	case <-resultCh:
	case <-time.After(100 * time.Millisecond):
	}
}

// TestModelQueueNodeDisconnectReleasesInflight verifies that a node disconnect
// event removes the node from queued item candidate lists, and that subsequent
// dispatch goes to a remaining live candidate (not the disconnected node).
func TestModelQueueNodeDisconnectReleasesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd2",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry1 := &edgenode.NodeEntry{NodeID: "node-nd1"}
	entry2 := &edgenode.NodeEntry{NodeID: "node-nd2"}
	// Queued item that can use either node (both at capacity 1).
	bothCands := []candidateNode{
		{entry: entry1, capacity: 1},
		{entry: entry2, capacity: 1},
	}

	m := newModelQueueManager(store)

	// Fill node-nd1 and node-nd2 capacities and record inflight.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-nd", groupPolicy{})
	g.inflight["node-nd1"] = 1
	g.inflight["node-nd2"] = 1
	// Queue an item that can use either node.
	item := &queueItem{
		candidates: bothCands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	g.queue = []*queueItem{item}
	m.mu.Unlock()
	trackTestLease(m, "g-nd", "run-nd-x", "node-nd1", "", false)
	trackTestLease(m, "g-nd", "run-nd-y", "node-nd2", "", false)

	// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
	// node-nd2 is still full, so no dispatch yet. Generation 0 fences the node
	// unconditionally (untracked/whole-node release).
	m.releaseNode("node-nd1", 0, "disconnected")

	// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
	m.releaseRun("run-nd-y", "complete")

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after disconnect+terminal, got error: %v", res.err)
		}
		if res.candidate == nil {
			t.Fatal("expected non-nil candidate after disconnect release")
		}
		// Must be nd2 — nd1 was removed from candidates on disconnect.
		if res.candidate.entry.NodeID != "node-nd2" {
			t.Errorf("expected dispatch to node-nd2, got %q (nd1 was disconnected)", res.candidate.entry.NodeID)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: queued item not dispatched after node disconnect and terminal event")
	}

	// The lease for run-nd-x should be gone (released at disconnect).
	m.mu.Lock()
	_, stillInFlight := m.leaseByRun["run-nd-x"]
	m.mu.Unlock()
	if stillInFlight {
		t.Error("run-nd-x lease should be released after node disconnect")
	}
}

// TestRunTerminalReleasesLeaseWhenEventSubscriberFull pins that lease accounting
// survives a lossy observability fanout. The bus drops into full subscriber
// channels by design, so the transport calls the service lifecycle hook directly;
// this test saturates the bus, confirms it is actually dropping, and then checks
// the terminal event still returns the slot.
func TestRunTerminalReleasesLeaseWhenEventSubscriberFull(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-drop",
		Providers: []config.NodeProviderConf{
			{ID: "prov-drop", Capacity: 1, Adapter: "vllm"},
		},
	})

	bus := edgeevents.NewBus()
	svc := New(edgenode.NewRegistry(), bus)
	svc.SetNodeStore(store)

	// A slow observer with a one-deep buffer: the fanout starts dropping as soon
	// as it stops reading, which is exactly the delivery correctness must not
	// depend on.
	_, unsub := bus.SubscribeAllRuns(1)
	defer unsub()

	entry := &edgenode.NodeEntry{NodeID: "node-drop"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-drop"}}
	admitted, _, err := svc.queue.admitWithReason(context.Background(), "g-drop", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}
	svc.queue.trackLease(admitted.leaseID, "run-drop-1")

	for i := 0; i < 8; i++ {
		bus.PublishRun(&iop.RunEvent{RunId: "run-drop-1", Type: "chunk"})
	}
	if dropped := bus.Stats().DroppedRunEvents; dropped == 0 {
		t.Fatal("expected the saturated subscriber to drop run events; the test is not exercising the drop path")
	}

	// The terminal event is delivered the way the transport delivers it.
	svc.HandleRunLifecycleEvent(&iop.RunEvent{RunId: "run-drop-1", Type: "complete"})

	svc.queue.mu.Lock()
	inflight, _ := svc.queue.getStatsForProviderLocked("node-drop", "prov-drop")
	svc.queue.mu.Unlock()
	if inflight != 0 {
		t.Errorf("in-flight after terminal event with a dropping bus = %d, want 0", inflight)
	}
	if got := leaseCount(svc.queue); got != 0 {
		t.Errorf("live leases after terminal event = %d, want 0", got)
	}
}

// TestModelQueueCancelledTerminalReleasesProviderInflight verifies that a
// cancelled terminal RunEvent releases provider in-flight accounting back to
// zero, the same as complete/error. This guards the BUG-2 recovery contract:
// once Edge OpenAI cancel propagation reaches Node and Node replies with a
// cancelled terminal event, the provider snapshot's in_flight must not stay
// stuck positive.
func TestModelQueueCancelledTerminalReleasesProviderInflight(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-x"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-1"}}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "node-x",
		Providers: []config.NodeProviderConf{
			{ID: "prov-1", Capacity: 1, Adapter: "vllm"},
		},
	})
	m := newModelQueueManager(store)
	m.setStore(store)

	selected, err := m.admit(context.Background(), "g-cancel", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || selected == nil {
		t.Fatalf("admit failed: %v", err)
	}

	m.trackLease(selected.leaseID, "run-cancel-1")

	m.mu.Lock()
	inf, _ := m.getStatsForProviderLocked("node-x", "prov-1")
	m.mu.Unlock()
	if inf != 1 {
		t.Fatalf("inflight before cancel: got %d, want 1", inf)
	}

	m.releaseRun("run-cancel-1", "cancelled")

	deadline := time.Now().Add(200 * time.Millisecond)
	for {
		m.mu.Lock()
		inf, _ = m.getStatsForProviderLocked("node-x", "prov-1")
		m.mu.Unlock()
		if inf == 0 {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("inflight did not recover to 0 after cancelled terminal event, got %d", inf)
		}
		time.Sleep(1 * time.Millisecond)
	}
}

// TestModelQueueContextCancelRemovesQueuedItem verifies that cancelling the
// context of a queued admit removes the item and returns context.Canceled.
func TestModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-cc1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// Fill the only slot.
	n, err := m.admit(context.Background(), "g-cc", "", "", cands, defPolicy, nil, false, false)
	if err != nil || n == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Queue a second admit with a cancellable context.
	ctx, cancel := context.WithCancel(context.Background())
	resultCh := make(chan error, 1)
	go func() {
		_, err := m.admit(ctx, "g-cc", "", "", cands, defPolicy, nil, false, false)
		resultCh <- err
	}()

	waitForQueueLen(t, m, "g-cc", 1)
	cancel()

	select {
	case err := <-resultCh:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("expected context.Canceled, got: %v", err)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: context cancellation not handled")
	}

	// Queue must be empty after cancellation.
	m.mu.Lock()
	qLen := 0
	if g, ok := m.groups["g-cc"]; ok {
		qLen = len(g.queue)
	}
	m.mu.Unlock()
	if qLen != 0 {
		t.Errorf("queue should be empty after cancel, got %d items", qLen)
	}
}

func TestProviderResourceCapacitySharedAcrossOrnithModelAliases(t *testing.T) {
	const (
		nodeID     = "rtx5090-lemonade-node"
		providerID = "rtx5090-lemonade"
	)
	modelAliases := []string{"ornith:35b", "ornith-fast"}
	entry := &edgenode.NodeEntry{NodeID: nodeID}
	cands := []candidateNode{
		{entry: entry, capacity: 1, providerID: providerID},
	}
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: nodeID,
		Providers: []config.NodeProviderConf{
			{ID: providerID, Capacity: 1, Adapter: "vllm"},
		},
	})
	m := newModelQueueManager(store)
	m.setStore(store)

	start := make(chan struct{})
	resCh := make(chan concurrentAdmitResult, 2)

	for i := 0; i < 2; i++ {
		go func(id int) {
			<-start
			ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
			defer cancel()
			groupKey := modelAliases[id]
			sel, err := m.admit(ctx, groupKey, "", "", cands, groupPolicy{maxQueue: 1, queueTimeout: 30 * time.Millisecond}, nil, false, false)
			resCh <- concurrentAdmitResult{candidate: sel, err: err}
		}(i)
	}

	close(start)

	var results []concurrentAdmitResult
	for i := 0; i < 2; i++ {
		results = append(results, <-resCh)
	}

	successCount := 0
	failCount := 0
	var successSel *candidateNode
	for _, res := range results {
		if res.err == nil {
			successCount++
			successSel = res.candidate
		} else {
			failCount++
		}
	}

	if successCount != 1 || failCount != 1 {
		t.Fatalf("expected exactly 1 success and 1 failure under concurrent admission, got success=%d fail=%d", successCount, failCount)
	}

	// Verify in-flight count peak is 1
	m.mu.Lock()
	res := m.resources[providerResourceKey{nodeID: nodeID, providerID: providerID}]
	if res == nil {
		m.mu.Unlock()
		t.Fatal("expected provider resource state to exist")
	}
	if res.inFlight != 1 {
		t.Errorf("expected inFlight peak to be 1, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Release the lease held by whichever alias won the shared provider slot.
	m.releaseLease(successSel.leaseID, "test-complete")

	// Verify count is 0
	m.mu.Lock()
	if res.inFlight != 0 {
		t.Errorf("expected inFlight to recover to 0, got %d", res.inFlight)
	}
	m.mu.Unlock()
}

func TestProviderSnapshotReadsAdmissionResourceState(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-snap-read", Alias: "node-snap-read"})

	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID:      "node-snap-read",
		Runtime: config.RuntimeConf{Concurrency: 5},
		Providers: []config.NodeProviderConf{
			{
				ID:                  "snap-prov-read",
				Adapter:             "vllm",
				Capacity:            3,
				LongContextCapacity: 2,
				Enabled:             nil,
			},
		},
	}
	store.Add(rec)

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)

	cands := []candidateNode{
		{entry: reg.All()[0], capacity: 3, providerID: "snap-prov-read", longContextCapacity: 2},
	}

	sel, err := svc.queue.admit(context.Background(), "g-snap-read", "", "", cands, groupPolicy{}, nil, true, false)
	if err != nil || sel == nil {
		t.Fatalf("admit failed: %v", err)
	}

	snaps := svc.ListNodeSnapshots()
	if len(snaps) != 1 {
		t.Fatalf("expected 1 node snapshot")
	}
	p := snaps[0].ProviderSnapshots[0]
	if p.InFlight != 1 {
		t.Errorf("expected InFlight 1, got %d", p.InFlight)
	}
	if p.LongInFlight != 1 {
		t.Errorf("expected LongInFlight 1, got %d", p.LongInFlight)
	}

	svc.queue.releaseSlotWithLong("g-snap-read", "node-snap-read", "snap-prov-read", true)

	snaps2 := svc.ListNodeSnapshots()
	p2 := snaps2[0].ProviderSnapshots[0]
	if p2.InFlight != 0 {
		t.Errorf("expected InFlight 0 after release, got %d", p2.InFlight)
	}
	if p2.LongInFlight != 0 {
		t.Errorf("expected LongInFlight 0 after release, got %d", p2.LongInFlight)
	}
}

func TestProviderResourceCapacityRefreshPreservesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-refresh",
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-refresh",
				Capacity: 2,
				Enabled:  nil, // defaults to true
			},
		},
	}
	store.Add(rec)

	m := newModelQueueManager(store)
	m.setStore(store)
	entry := &edgenode.NodeEntry{NodeID: "node-refresh"}
	cands := []candidateNode{
		{entry: entry, capacity: 2, providerID: "prov-refresh"},
	}

	// Admit 2 items (matching capacity)
	sel1, err := m.admit(context.Background(), "group-1", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || sel1 == nil {
		t.Fatalf("admit 1 failed: %v", err)
	}
	sel2, err := m.admit(context.Background(), "group-2", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || sel2 == nil {
		t.Fatalf("admit 2 failed: %v", err)
	}

	// Verify inFlight = 2
	m.mu.Lock()
	res := m.resources[providerResourceKey{nodeID: "node-refresh", providerID: "prov-refresh"}]
	if res.inFlight != 2 {
		t.Fatalf("expected inFlight = 2, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Update store with decreased capacity (2 -> 1)
	newStore := edgenode.NewNodeStore()
	newRec := &edgenode.NodeRecord{
		ID: "node-refresh",
		Providers: []config.NodeProviderConf{
			{
				ID:       "prov-refresh",
				Capacity: 1,
				Enabled:  nil,
			},
		},
	}
	newStore.Add(newRec)

	// Refresh store
	m.setStore(newStore)

	// Verify capacity is updated but inFlight is preserved
	m.mu.Lock()
	if res.capacity != 1 {
		t.Errorf("expected updated capacity to be 1, got %d", res.capacity)
	}
	if res.inFlight != 2 {
		t.Errorf("expected inFlight to remain 2, got %d", res.inFlight)
	}
	m.mu.Unlock()

	// Third admit should fail because capacity is full (capacity 1, inflight 2)
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	defer cancel()
	_, err = m.admit(ctx, "group-3", "", "", cands, groupPolicy{queueTimeout: 10 * time.Millisecond}, nil, false, false)
	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected queue timeout for third admit, got: %v", err)
	}
}

// TestProviderResourceRemovedProviderDoesNotReadmitQueuedCandidate verifies
// that a stale queued candidate cannot re-admit a removed provider. After the
// last release deletes the orphan resource, a queue pump must not recreate an
// active resource from the stale candidate's old capacity.
//
// Arrange order is critical: the second admit must be in-flight in the queue
// BEFORE the release that triggers pumpAllLocked. This exercises the real
// release-triggered queue-pump path where releaseSlotLocked → pumpAllLocked
// would otherwise hand a stale candidate to a deleted provider.
func TestProviderResourceRemovedProviderDoesNotReadmitQueuedCandidate(t *testing.T) {
	store := edgenode.NewNodeStore()
	rec := &edgenode.NodeRecord{
		ID: "node-stale",
		Providers: []config.NodeProviderConf{
			{ID: "prov-stale", Capacity: 1, Adapter: "vllm"},
		},
	}
	store.Add(rec)

	entry := &edgenode.NodeEntry{NodeID: "node-stale"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-stale"}}

	m := newModelQueueManager(store)
	m.setStore(store)

	// 1. Admit the first request — fills the only capacity slot.
	first, err := m.admit(context.Background(), "g-stale", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	key := providerResourceKey{nodeID: "node-stale", providerID: "prov-stale"}
	m.mu.Lock()
	res := m.resources[key]
	m.mu.Unlock()
	if res == nil || res.orphan {
		t.Fatal("expected active non-orphan resource after first admit")
	}

	// 2. Start the second admit in a goroutine so it blocks in the queue
	//    while the first request is still running. Use a short timeout so the
	//    test can observe the terminal result quickly.
	resultCh := make(chan admitResult, 1)
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
		defer cancel()
		c, _, e := m.admitWithReason(ctx, "g-stale", "", "", cands, groupPolicy{maxQueue: 16, queueTimeout: 500 * time.Millisecond}, nil, false, false)
		resultCh <- admitResult{candidate: c, err: e}
	}()

	// 3. Confirm the second request actually entered the queue.
	waitForQueueLen(t, m, "g-stale", 1)

	// 4. Remove the provider from the store so the resource becomes orphan.
	rec.Providers = nil
	m.setStore(store)

	m.mu.Lock()
	res, ok := m.resources[key]
	m.mu.Unlock()
	if !ok || res == nil || !res.orphan {
		t.Fatalf("expected resource to be orphan after store update, ok=%t, res=%v", ok, res)
	}

	// 5. Release the first slot to trigger the queue pump. The orphan
	//    resource must be deleted (inFlight=0) and the queue pump must
	//    NOT recreate it via pumpAllLocked → findAvailableNodeLocked.
	m.releaseSlotWithLong("g-stale", "node-stale", "prov-stale", false)

	// 6. The second request must NOT succeed — the stale candidate must
	//    not be dispatched and the resource must not be recreated.
	select {
	case res := <-resultCh:
		if res.err == nil {
			t.Fatal("expected error/timeout for queued candidate of removed provider, got nil error")
		}
		if res.candidate != nil {
			t.Fatalf("expected no candidate dispatch, got %+v", res.candidate)
		}
	case <-time.After(1 * time.Second):
		t.Fatal("timeout: second admit did not return within expected window")
	}

	// 7. Verify the resource was not recreated as an active entry.
	m.mu.Lock()
	_, exists := m.resources[key]
	m.mu.Unlock()
	if exists {
		t.Fatal("expected orphan resource to remain deleted after queue pump")
	}
}

// queueItemForTest builds a pending item with a generous deadline so the test
// controls when it resolves.
func queueItemForTest(cands []candidateNode, long bool) *queueItem {
	return &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(5 * time.Second),
		long:       long,
	}
}

// enqueueForTest registers an item in the given group through the same helper
// admission uses, so it receives a real global enqueue sequence. The optional
// resolveCandidates closure is forwarded to the manager; nil keeps the legacy
// snapshot path for existing fixtures.
func enqueueForTest(m *modelQueueManager, groupKey string, item *queueItem, resolveCandidates func() ([]candidateNode, error)) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.enqueueItemLocked(m.getOrCreateGroupLocked(groupKey, groupPolicy{}), item, resolveCandidates)
}

// TestGlobalPumpWakesDifferentModelGroup verifies the cross-group wake-up
// contract (S03 / cross-group-wakeup): provider resources are shared across
// model groups, so a lease released by one group must wake a waiter queued in
// another. Before the global pump, the waiter slept until its own group saw an
// event or its queue timeout expired.
func TestGlobalPumpWakesDifferentModelGroup(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-wake"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-wake"}}
	m := newModelQueueManager(nil)

	// Group A takes the provider's only slot.
	first, err := m.admit(context.Background(), "g-wake-a", "", "", cands, groupPolicy{}, nil, false, false)
	if err != nil || first == nil {
		t.Fatalf("group A admit: %v", err)
	}

	// Group B queues a request for the same provider.
	waiter := queueItemForTest(cands, false)
	enqueueForTest(m, "g-wake-b", waiter, nil)

	// Releasing group A's lease must dispatch group B's waiter.
	m.releaseLease(first.leaseID, "complete")

	select {
	case res := <-waiter.waitCh:
		if res.err != nil {
			t.Fatalf("group B waiter expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.providerID != "prov-wake" {
			t.Fatalf("group B waiter: unexpected candidate %v", res.candidate)
		}
	default:
		t.Fatal("group B waiter was not woken by the group A release (cross-group wake-up missing)")
	}
}

// TestGlobalPumpPreservesEarliestDispatchableSequence verifies that the pump
// orders waiters by global arrival, not by group: an item blocked on a full
// provider is skipped, and among the items that can dispatch the earliest
// enqueued one wins even when it sits in a different group than the release.
func TestGlobalPumpPreservesEarliestDispatchableSequence(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-seq"}
	blockedCands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-seq-blocked"}}
	openCands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-seq-open"}}
	m := newModelQueueManager(nil)

	// Fill both providers so every queued item starts out blocked.
	if _, err := m.admit(context.Background(), "g-seq-a", "", "", blockedCands, groupPolicy{}, nil, false, false); err != nil {
		t.Fatalf("fill blocked provider: %v", err)
	}
	openHolder, err := m.admit(context.Background(), "g-seq-a", "", "", openCands, groupPolicy{}, nil, false, false)
	if err != nil {
		t.Fatalf("fill open provider: %v", err)
	}

	// Arrival order: the earliest item can never dispatch (its provider is never
	// released); the next two compete for the one slot that does free.
	blocked := queueItemForTest(blockedCands, false)
	enqueueForTest(m, "g-seq-a", blocked, nil)
	earlier := queueItemForTest(openCands, false)
	enqueueForTest(m, "g-seq-b", earlier, nil)
	later := queueItemForTest(openCands, false)
	enqueueForTest(m, "g-seq-a", later, nil)

	m.releaseLease(openHolder.leaseID, "complete")

	select {
	case res := <-earlier.waitCh:
		if res.err != nil {
			t.Fatalf("earlier item dispatch error: %v", res.err)
		}
	default:
		t.Fatal("earliest dispatchable item was not dispatched first")
	}
	select {
	case <-later.waitCh:
		t.Fatal("later item dispatched out of global enqueue order")
	default:
	}
	select {
	case <-blocked.waitCh:
		t.Fatal("item on a still-full provider must stay queued")
	default:
	}

	// Freeing the slot again promotes the later item; the blocked one stays put,
	// proving skipping does not consume its turn.
	m.mu.Lock()
	pending := len(m.pendingItemsLocked())
	m.mu.Unlock()
	if pending != 2 {
		t.Fatalf("expected 2 pending items after the first dispatch, got %d", pending)
	}
}

// TestGlobalPumpConcurrentReleaseCancelRefreshNoLeak drives releases, cancels,
// and config refreshes against the global pump concurrently. The invariants it
// pins are the ones a shared scheduler can silently break: no request is handed
// two leases, provider in-flight never exceeds configured capacity, and every
// counter returns to zero once the storm ends.
func TestGlobalPumpConcurrentReleaseCancelRefreshNoLeak(t *testing.T) {
	const (
		capacity   = 3
		groups     = 3
		perGroup   = 12
		nodeID     = "node-race"
		providerID = "prov-race"
	)

	newStore := func(cap int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: nodeID,
			Providers: []config.NodeProviderConf{
				{ID: providerID, Adapter: "vllm", Capacity: cap, LongContextCapacity: 1},
			},
		})
		return store
	}

	m := newModelQueueManager(nil)
	m.setStore(newStore(capacity))

	entry := &edgenode.NodeEntry{NodeID: nodeID}
	cands := []candidateNode{{entry: entry, capacity: capacity, providerID: providerID, longContextCapacity: 1}}

	var (
		mu         sync.Mutex
		seenLeases = map[uint64]int{}
		peak       int
		peakLong   int
	)
	observe := func() {
		inFlight, longInFlight := providerResourceCounts(m, nodeID, providerID)
		mu.Lock()
		if inFlight > peak {
			peak = inFlight
		}
		if longInFlight > peakLong {
			peakLong = longInFlight
		}
		mu.Unlock()
	}

	stop := make(chan struct{})
	var refreshWG sync.WaitGroup
	refreshWG.Add(1)
	go func() {
		defer refreshWG.Done()
		for {
			select {
			case <-stop:
				return
			default:
			}
			// Config refresh alternates capacity; both values are pumped, so the
			// scheduler re-evaluates every waiter on each swap.
			m.setStore(newStore(capacity))
			m.setStore(newStore(capacity + 1))
			observe()
		}
	}()

	var wg sync.WaitGroup
	for g := 0; g < groups; g++ {
		groupKey := fmt.Sprintf("g-race-%d", g)
		for i := 0; i < perGroup; i++ {
			wg.Add(1)
			long := i%3 == 0
			// A share of the callers abandon their request mid-flight, so cancel
			// races the dispatch hand-off.
			cancelEarly := i%4 == 0
			go func() {
				defer wg.Done()
				ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
				defer cancel()
				if cancelEarly {
					go func() {
						time.Sleep(time.Millisecond)
						cancel()
					}()
				}
				admitted, err := m.admit(ctx, groupKey, "", "", cands, groupPolicy{maxQueue: 64, queueTimeout: 150 * time.Millisecond}, nil, long, false)
				if err != nil || admitted == nil {
					return
				}
				mu.Lock()
				seenLeases[admitted.leaseID]++
				mu.Unlock()
				observe()
				time.Sleep(time.Millisecond)
				m.releaseLease(admitted.leaseID, "complete")
			}()
		}
	}

	wg.Wait()
	close(stop)
	refreshWG.Wait()

	// Every admitted request holds exactly one lease identity.
	mu.Lock()
	for id, count := range seenLeases {
		if count != 1 {
			mu.Unlock()
			t.Fatalf("lease %d handed out %d times; a request must own exactly one lease", id, count)
		}
	}
	observedPeak, observedPeakLong := peak, peakLong
	mu.Unlock()

	if observedPeak > capacity+1 {
		t.Errorf("provider in-flight peaked at %d, above the highest configured capacity %d", observedPeak, capacity+1)
	}
	if observedPeakLong > 1 {
		t.Errorf("long in-flight peaked at %d, above the configured long capacity 1", observedPeakLong)
	}

	// Final state must be fully drained: no held resource, lease, or pending item.
	inFlight, longInFlight := providerResourceCounts(m, nodeID, providerID)
	if inFlight != 0 || longInFlight != 0 {
		t.Errorf("expected zero in-flight after drain, got inflight=%d long=%d", inFlight, longInFlight)
	}
	if got := leaseCount(m); got != 0 {
		t.Errorf("expected zero live leases after drain, got %d", got)
	}
	if got := pendingItemCount(m); got != 0 {
		t.Errorf("expected zero pending items after drain, got %d", got)
	}
}

// TestModelQueueTerminalUnavailableOnNoCandidates verifies that when a queued
// item's live resolver returns zero candidates, the item is removed from the
// queue and the waiter receives a typed errProviderUnavailable error exactly
// once. This is the terminal-unavailable outcome contract.
func TestModelQueueTerminalUnavailableOnNoCandidates(t *testing.T) {
	store := edgenode.NewNodeStore()
	m := newModelQueueManager(store)

	candidates := []candidateNode{
		{entry: &edgenode.NodeEntry{NodeID: "node-x"}, capacity: 1, providerID: "prov-x"},
	}
	item := &queueItem{
		candidates: candidates,
		resolveCandidates: func() ([]candidateNode, error) {
			return nil, nil // no live candidates
		},
		waitCh:       make(chan admitResult, 1),
		deadline:     time.Now().Add(5 * time.Second),
		providerPool: true,
	}
	enqueueForTest(m, "g-no-cand", item, item.resolveCandidates)

	// Run the pump — should deliver terminal unavailable and remove the item.
	m.mu.Lock()
	m.pumpAllLocked()
	m.mu.Unlock()

	// The waiter should have received the terminal unavailable error.
	select {
	case res := <-item.waitCh:
		if res.candidate != nil {
			t.Fatalf("expected no candidate, got %+v", res.candidate)
		}
		if !errors.Is(res.err, errProviderUnavailable) {
			t.Fatalf("expected errProviderUnavailable, got: %v", res.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: waiter was not notified of terminal unavailable")
	}

	// Queue must be empty after terminal delivery.
	m.mu.Lock()
	qLen := len(m.groups["g-no-cand"].queue)
	m.mu.Unlock()
	if qLen != 0 {
		t.Errorf("expected queue to be empty after terminal delivery, got %d", qLen)
	}
}

type providerPoolQueueFixture struct {
	svc        *Service
	registry   *edgenode.Registry
	store      *edgenode.NodeStore
	catalog    []config.ModelCatalogEntry
	policy     groupPolicy
	req        SubmitRunRequest
	candidates []candidateNode
	resolver   func() ([]candidateNode, error)
	entry      *edgenode.NodeEntry
}

func newProviderPoolQueueFixture(t *testing.T, suffix string) *providerPoolQueueFixture {
	t.Helper()
	nodeID := "node-" + suffix
	providerID := "provider-" + suffix
	servedModel := "served-" + suffix
	groupKey := "group-" + suffix

	registry := edgenode.NewRegistry()
	entry := &edgenode.NodeEntry{NodeID: nodeID, Alias: nodeID}
	registry.Register(entry)
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: nodeID,
		Providers: []config.NodeProviderConf{
			{
				ID:       providerID,
				Type:     "vllm",
				Category: config.CategoryAPI,
				Models:   []string{servedModel},
				Health:   "available",
				Capacity: 1,
			},
		},
	})
	catalog := []config.ModelCatalogEntry{
		{ID: groupKey, Providers: map[string]string{providerID: servedModel}},
	}
	policy := NewGroupPolicy(16, 5*time.Second)
	svc := New(registry, edgeevents.NewBus())
	svc.SetRuntimeConfig(store, catalog, policy)
	req := SubmitRunRequest{ModelGroupKey: groupKey, ProviderPool: true}
	candidates, _, err := svc.resolveQueueCandidates(req)
	if err != nil {
		t.Fatalf("resolve initial provider-pool candidates: %v", err)
	}

	return &providerPoolQueueFixture{
		svc:        svc,
		registry:   registry,
		store:      store,
		catalog:    catalog,
		policy:     policy,
		req:        req,
		candidates: candidates,
		resolver:   svc.resolveQueueCandidatesClosure(req),
		entry:      entry,
	}
}

func (f *providerPoolQueueFixture) admitHolder(t *testing.T) *candidateNode {
	t.Helper()
	holder, err := f.svc.queue.admit(
		t.Context(), f.req.ModelGroupKey, "", "", f.candidates, f.policy,
		f.resolver, false, true,
	)
	if err != nil {
		t.Fatalf("admit capacity holder: %v", err)
	}
	f.svc.queue.trackLease(holder.leaseID, "run-holder-"+f.req.ModelGroupKey)
	return holder
}

func (f *providerPoolQueueFixture) enqueueWaiter() *queueItem {
	item := queueItemForTest(f.candidates, false)
	item.providerPool = true
	enqueueForTest(f.svc.queue, f.req.ModelGroupKey, item, f.resolver)
	return item
}

func providerQueuePressureForTest(m *modelQueueManager, nodeID, providerID string) (queued, longQueued int) {
	m.mu.Lock()
	defer m.mu.Unlock()
	pressure := m.providerQueuePressureLocked()
	if byProvider := pressure[nodeID]; byProvider != nil {
		if value := byProvider[providerID]; value != nil {
			return value.queued, value.longQueued
		}
	}
	return 0, 0
}

// TestProviderPoolQueuedDisconnectReturnsUnavailableImmediately verifies that
// the production registry/store/catalog resolver classifies the last live
// provider's disconnect as terminal unavailable. The queue and candidate
// pressure are cleared in the same disconnect pump, and a later pump cannot
// deliver a duplicate terminal result.
func TestProviderPoolQueuedDisconnectReturnsUnavailableImmediately(t *testing.T) {
	fixture := newProviderPoolQueueFixture(t, "disconnect-unavailable")
	holder := fixture.admitHolder(t)
	item := fixture.enqueueWaiter()

	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending before disconnect = %d, want 1", pending)
	}
	queued, longQueued := providerQueuePressureForTest(
		fixture.svc.queue, holder.entry.NodeID, holder.providerID,
	)
	if queued != 1 || longQueued != 0 {
		t.Fatalf("candidate pressure before disconnect = (%d, %d), want (1, 0)", queued, longQueued)
	}

	// Keep an unrelated Node connected so production resolution reaches the
	// filtered zero-candidate branch instead of the simpler empty-registry branch.
	fixture.registry.Register(&edgenode.NodeEntry{NodeID: "node-unrelated-disconnect-unavailable"})
	fixture.registry.Unregister(fixture.entry.NodeID)
	fixture.svc.HandleNodeDisconnect(
		fixture.entry.NodeID, fixture.entry.ConnectionGeneration, "test disconnect",
	)

	select {
	case result := <-item.waitCh:
		if result.candidate != nil {
			t.Fatalf("disconnect waiter candidate = %+v, want nil", result.candidate)
		}
		if !errors.Is(result.err, errProviderUnavailable) {
			t.Fatalf("disconnect waiter error = %v, want errProviderUnavailable", result.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("disconnect waiter did not receive terminal unavailable immediately")
	}

	if pending := pendingItemCount(fixture.svc.queue); pending != 0 {
		t.Fatalf("pending after disconnect = %d, want 0", pending)
	}
	queued, longQueued = providerQueuePressureForTest(
		fixture.svc.queue, holder.entry.NodeID, holder.providerID,
	)
	if queued != 0 || longQueued != 0 {
		t.Fatalf("candidate pressure after disconnect = (%d, %d), want (0, 0)", queued, longQueued)
	}
	if inFlight, _ := providerResourceCounts(fixture.svc.queue, holder.entry.NodeID, holder.providerID); inFlight != 0 {
		t.Fatalf("in-flight after disconnect = %d, want 0", inFlight)
	}
	if leases := leaseCount(fixture.svc.queue); leases != 0 {
		t.Fatalf("live leases after disconnect = %d, want 0", leases)
	}

	fixture.svc.queue.mu.Lock()
	fixture.svc.queue.pumpAllLocked()
	fixture.svc.queue.mu.Unlock()
	select {
	case duplicate := <-item.waitCh:
		t.Fatalf("duplicate terminal delivery after re-pump: %+v", duplicate)
	default:
	}
}

// TestModelQueueStaysQueuedWhenAllCandidatesFull verifies through the
// production provider-pool resolver that a connected, available provider at
// capacity is a temporary block, not terminal unavailability.
func TestModelQueueStaysQueuedWhenAllCandidatesFull(t *testing.T) {
	fixture := newProviderPoolQueueFixture(t, "all-candidates-full")
	holder := fixture.admitHolder(t)
	item := fixture.enqueueWaiter()

	// Pump — the item should be skipped (temporarily blocked), not delivered.
	fixture.svc.queue.mu.Lock()
	fixture.svc.queue.pumpAllLocked()
	fixture.svc.queue.mu.Unlock()

	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending when all candidates full = %d, want 1", pending)
	}
	queued, longQueued := providerQueuePressureForTest(
		fixture.svc.queue, holder.entry.NodeID, holder.providerID,
	)
	if queued != 1 || longQueued != 0 {
		t.Fatalf("candidate pressure while full = (%d, %d), want (1, 0)", queued, longQueued)
	}
	select {
	case <-item.waitCh:
		t.Fatal("waiter should not have been notified when all candidates full")
	default:
	}

	fixture.svc.HandleRunLifecycleEvent(&iop.RunEvent{
		RunId: "run-holder-" + fixture.req.ModelGroupKey,
		Type:  "complete",
	})
	select {
	case result := <-item.waitCh:
		if result.err != nil || result.candidate == nil {
			t.Fatalf("waiter after holder release = %+v, want dispatch", result)
		}
		fixture.svc.queue.releaseLease(result.candidate.leaseID, "test cleanup")
	case <-time.After(100 * time.Millisecond):
		t.Fatal("full candidate waiter did not dispatch after holder release")
	}
}

// TestProviderPoolQueuedResolverErrorStaysQueuedAndRecovers verifies that a
// production catalog/config resolver fault is not confused with the typed
// no-live-provider outcome. The waiter remains queued across the fault and
// dispatches after a valid snapshot is restored and capacity is released.
func TestProviderPoolQueuedResolverErrorStaysQueuedAndRecovers(t *testing.T) {
	fixture := newProviderPoolQueueFixture(t, "resolver-recovery")
	fixture.admitHolder(t)
	item := fixture.enqueueWaiter()

	// Remove the model catalog entry while the request is queued. This is a
	// resolver/config fault, not a live-candidate absence, so it must not carry
	// errProviderUnavailable or terminate the waiter.
	fixture.svc.SetRuntimeConfig(fixture.store, nil, fixture.policy)
	if _, err := fixture.resolver(); err == nil {
		t.Fatal("resolver with missing catalog returned nil error")
	} else if errors.Is(err, errProviderUnavailable) {
		t.Fatalf("missing catalog error incorrectly classified unavailable: %v", err)
	}
	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending during resolver fault = %d, want 1", pending)
	}
	select {
	case result := <-item.waitCh:
		t.Fatalf("resolver fault unexpectedly notified waiter: %+v", result)
	default:
	}

	fixture.svc.SetRuntimeConfig(fixture.store, fixture.catalog, fixture.policy)
	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending after resolver recovery while full = %d, want 1", pending)
	}
	queued, longQueued := providerQueuePressureForTest(
		fixture.svc.queue, fixture.entry.NodeID, fixture.candidates[0].providerID,
	)
	if queued != 1 || longQueued != 0 {
		t.Fatalf("candidate pressure after resolver recovery = (%d, %d), want (1, 0)", queued, longQueued)
	}

	fixture.svc.HandleRunLifecycleEvent(&iop.RunEvent{
		RunId: "run-holder-" + fixture.req.ModelGroupKey,
		Type:  "complete",
	})
	select {
	case result := <-item.waitCh:
		if result.err != nil || result.candidate == nil {
			t.Fatalf("waiter after resolver recovery = %+v, want dispatch", result)
		}
		fixture.svc.queue.releaseLease(result.candidate.leaseID, "test cleanup")
	case <-time.After(100 * time.Millisecond):
		t.Fatal("resolver-recovered waiter did not dispatch after capacity release")
	}
}

// TestProviderPoolDisconnectAuthoritativeSettlementNoLeak verifies that an
// authoritative disconnect with the live resolver settles the last live
// candidate as terminal unavailable in a single deterministic pass: the waiter
// receives errProviderUnavailable, the queue is empty, in-flight and candidate
// pressure are zero, and every lease is released. It also runs with a
// saturated event subscriber to prove that event bus backpressure does not
// affect the settlement contract — the transport drives releaseNode directly.
func TestProviderPoolDisconnectAuthoritativeSettlementNoLeak(t *testing.T) {
	bus := edgeevents.NewBus()
	fixture := newProviderPoolQueueFixture(t, "auth-settlement")
	// Override the bus so we can verify it is actually saturated (dropping).
	// We achieve this by subscribing with a one-deep buffer and flooding it
	// before the disconnect.
	fixture.svc = New(fixture.registry, bus)
	fixture.svc.SetRuntimeConfig(fixture.store, fixture.catalog, fixture.policy)
	// Re-resolve candidates against the new service so the resolver closure
	// points at the correct service instance.
	cands, _, err := fixture.svc.resolveQueueCandidates(fixture.req)
	if err != nil {
		t.Fatalf("re-resolve candidates after bus swap: %v", err)
	}
	fixture.candidates = cands
	fixture.resolver = fixture.svc.resolveQueueCandidatesClosure(fixture.req)

	// Saturated subscriber: one-deep buffer, then flood so drops are real.
	_, unsub := bus.SubscribeAllRuns(1)
	defer unsub()
	for i := 0; i < 8; i++ {
		bus.PublishRun(&iop.RunEvent{RunId: "run-flood", Type: "chunk"})
	}
	if dropped := bus.Stats().DroppedRunEvents; dropped == 0 {
		t.Fatal("expected the saturated subscriber to drop events; the test must exercise the drop path")
	}

	holder := fixture.admitHolder(t)
	item := fixture.enqueueWaiter()

	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending before disconnect = %d, want 1", pending)
	}

	// Keep an unrelated node connected so production resolution reaches the
	// filtered zero-candidate branch (not the empty-registry shortcut).
	fixture.registry.Register(&edgenode.NodeEntry{NodeID: "node-unrelated-auth"})

	// Authoritative disconnect: the registry confirms the disconnecting
	// client still owned the entry.
	fixture.svc.HandleNodeDisconnect(
		fixture.entry.NodeID, fixture.entry.ConnectionGeneration, "test disconnect",
	)

	// Waiter must receive typed unavailable exactly once.
	select {
	case result := <-item.waitCh:
		if result.candidate != nil {
			t.Fatalf("disconnect waiter candidate = %+v, want nil", result.candidate)
		}
		if !errors.Is(result.err, errProviderUnavailable) {
			t.Fatalf("disconnect waiter error = %v, want errProviderUnavailable", result.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("disconnect waiter did not receive terminal unavailable immediately")
	}

	// Post-settlement invariants.
	if pending := pendingItemCount(fixture.svc.queue); pending != 0 {
		t.Fatalf("pending after authoritative disconnect = %d, want 0", pending)
	}
	queued, longQueued := providerQueuePressureForTest(
		fixture.svc.queue, holder.entry.NodeID, holder.providerID,
	)
	if queued != 0 || longQueued != 0 {
		t.Fatalf("candidate pressure after disconnect = (%d, %d), want (0, 0)", queued, longQueued)
	}
	if inFlight, _ := providerResourceCounts(fixture.svc.queue, holder.entry.NodeID, holder.providerID); inFlight != 0 {
		t.Fatalf("in-flight after disconnect = %d, want 0", inFlight)
	}
	if leases := leaseCount(fixture.svc.queue); leases != 0 {
		t.Fatalf("live leases after disconnect = %d, want 0", leases)
	}

	// Re-pump must not deliver a duplicate terminal result.
	fixture.svc.queue.mu.Lock()
	fixture.svc.queue.pumpAllLocked()
	fixture.svc.queue.mu.Unlock()
	select {
	case duplicate := <-item.waitCh:
		t.Fatalf("duplicate terminal delivery after re-pump: %+v", duplicate)
	default:
	}
}

// TestProviderPoolLastProviderDisconnectTerminalSettlement verifies that when
// the very last live candidate for a queued provider-pool item disappears
// (because the node disconnects), the waiter receives errProviderUnavailable
// immediately — not via queue timeout — and all accounting (queue, lease,
// in-flight, candidate pressure) returns to zero in the same pass.
func TestProviderPoolLastProviderDisconnectTerminalSettlement(t *testing.T) {
	fixture := newProviderPoolQueueFixture(t, "last-provider")
	holder := fixture.admitHolder(t)
	item := fixture.enqueueWaiter()

	// Verify pre-disconnect state.
	if pending := pendingItemCount(fixture.svc.queue); pending != 1 {
		t.Fatalf("pending before last-provider disconnect = %d, want 1", pending)
	}

	// The last live candidate disappears via disconnect.
	fixture.svc.HandleNodeDisconnect(
		fixture.entry.NodeID, fixture.entry.ConnectionGeneration, "last provider disconnect",
	)

	select {
	case result := <-item.waitCh:
		if result.candidate != nil {
			t.Fatalf("last-provider waiter candidate = %+v, want nil", result.candidate)
		}
		if !errors.Is(result.err, errProviderUnavailable) {
			t.Fatalf("last-provider waiter error = %v, want errProviderUnavailable", result.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("last-provider waiter did not receive terminal unavailable immediately")
	}

	// Queue, lease, and pressure must all be zero.
	if pending := pendingItemCount(fixture.svc.queue); pending != 0 {
		t.Fatalf("pending after last-provider disconnect = %d, want 0", pending)
	}
	if leases := leaseCount(fixture.svc.queue); leases != 0 {
		t.Fatalf("live leases after last-provider disconnect = %d, want 0", leases)
	}

	// No duplicate terminal delivery after a subsequent pump.
	fixture.svc.queue.mu.Lock()
	fixture.svc.queue.pumpAllLocked()
	fixture.svc.queue.mu.Unlock()
	select {
	case duplicate := <-item.waitCh:
		t.Fatalf("duplicate terminal delivery after re-pump: %+v", duplicate)
	default:
	}

	// Also verify that the holder's resource is cleaned up.
	if inFlight, _ := providerResourceCounts(fixture.svc.queue, holder.entry.NodeID, holder.providerID); inFlight != 0 {
		t.Fatalf("holder in-flight after last-provider disconnect = %d, want 0", inFlight)
	}
}