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

package service

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

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

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

// The recovery-preference tests exercise the production admission path
// (admitWithRecovery → applyRecoveryPreferenceLocked → findAvailableNodeLocked →
// reserveCandidateLocked) rather than the pure helper. The avoided provider
// "prov-a-primary" sorts before the alternate "prov-b-backup", so the plain
// rotation would pick the avoided provider; a case that instead selects the
// alternate proves the preference actually changed the dispatched candidate.
const (
	recoveryNodeID    = "node-recovery"
	recoveryAvoidID   = "prov-a-primary"
	recoveryAltID     = "prov-b-backup"
	recoveryServed    = "served-x"
	recoveryGroupKey  = "recovery-model"
	recoveryAvoidAdap = "vllm-a"
	recoveryAltAdap   = "vllm-b"
)

// newRecoveryQueueFixture builds a store-backed queue manager with two
// capacity-1 providers on one node and the provider-pool policy seeded from the
// store, matching how production reconciles resources before admission.
func newRecoveryQueueFixture(t *testing.T) (*modelQueueManager, *edgenode.NodeEntry) {
	t.Helper()
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      recoveryNodeID,
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: recoveryAvoidID, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
			{ID: recoveryAltID, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
		},
	})
	m := newModelQueueManager(store)
	m.setProviderPoolPolicyLocked(store, NewGroupPolicy(4, 5*time.Second))
	entry := &edgenode.NodeEntry{NodeID: recoveryNodeID}
	return m, entry
}

func recoveryCandidate(entry *edgenode.NodeEntry, providerID string) candidateNode {
	return candidateNode{
		entry:        entry,
		providerID:   providerID,
		servedTarget: recoveryServed,
		capacity:     1,
		generation:   entry.ConnectionGeneration,
	}
}

// markRecoveryUnavailable installs a runtime-health overlay that lowers one
// provider on the candidate's generation, without mutating any config.
func markRecoveryUnavailable(m *modelQueueManager, entry *edgenode.NodeEntry, providerID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.runtimeHealth[providerRuntimeHealthKey{
		nodeID:     entry.NodeID,
		generation: entry.ConnectionGeneration,
		providerID: providerID,
	}] = &providerRuntimeHealthOverlay{
		adapter:        recoveryAvoidAdap,
		target:         recoveryServed,
		observationSeq: 1,
		unavailable:    true,
	}
}

// TestProviderRecoverySelectionImmediateAdmission drives the immediate
// provider-pool admission path for every recovery branch and asserts the
// selected provider (or typed rejection), exactly one lease per dispatch, and
// that every counter settles back to zero after release with no forbidden
// reservation on a rejected policy.
func TestProviderRecoverySelectionImmediateAdmission(t *testing.T) {
	cases := []struct {
		name           string
		candidates     []string // provider ids present in the request
		unavailable    string   // provider id lowered by runtime overlay, or ""
		recovery       recoveryCandidatePolicy
		wantProviderID string // expected dispatched provider, or "" when rejected
		wantErr        error  // expected terminal error, or nil on dispatch
	}{
		{
			name:           "eligible_alternate_preferred_over_avoided",
			candidates:     []string{recoveryAvoidID, recoveryAltID},
			recovery:       recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID},
			wantProviderID: recoveryAltID,
		},
		{
			name:           "same_only_fallback_true_selects_avoided",
			candidates:     []string{recoveryAvoidID},
			recovery:       recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID, allowAvoidedProviderFallback: true},
			wantProviderID: recoveryAvoidID,
		},
		{
			name:       "same_only_fallback_false_rejects",
			candidates: []string{recoveryAvoidID},
			recovery:   recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID},
			wantErr:    ErrProviderPoolCandidateRejected,
		},
		{
			name:           "unavailable_alternate_fallback_true_selects_avoided",
			candidates:     []string{recoveryAvoidID, recoveryAltID},
			unavailable:    recoveryAltID,
			recovery:       recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID, allowAvoidedProviderFallback: true},
			wantProviderID: recoveryAvoidID,
		},
		{
			name:        "unavailable_alternate_fallback_false_rejects",
			candidates:  []string{recoveryAvoidID, recoveryAltID},
			unavailable: recoveryAltID,
			recovery:    recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID},
			wantErr:     ErrProviderPoolCandidateRejected,
		},
		{
			name:           "unavailable_avoided_selects_alternate",
			candidates:     []string{recoveryAvoidID, recoveryAltID},
			unavailable:    recoveryAvoidID,
			recovery:       recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID},
			wantProviderID: recoveryAltID,
		},
		{
			name:        "same_only_runtime_unavailable_is_terminal",
			candidates:  []string{recoveryAvoidID},
			unavailable: recoveryAvoidID,
			recovery:    recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID, allowAvoidedProviderFallback: true},
			wantErr:     errProviderUnavailable,
		},
		{
			name:           "empty_hints_dispatches_by_rotation",
			candidates:     []string{recoveryAvoidID, recoveryAltID},
			recovery:       recoveryCandidatePolicy{},
			wantProviderID: recoveryAvoidID, // lowest providerID wins the rotation
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m, entry := newRecoveryQueueFixture(t)
			if tc.unavailable != "" {
				markRecoveryUnavailable(m, entry, tc.unavailable)
			}
			candidates := make([]candidateNode, 0, len(tc.candidates))
			for _, id := range tc.candidates {
				candidates = append(candidates, recoveryCandidate(entry, id))
			}

			ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
			defer cancel()
			selected, _, err := m.admitWithRecovery(ctx, recoveryGroupKey, "", recoveryServed, candidates, groupPolicy{}, nil, false, true, tc.recovery)

			if tc.wantErr != nil {
				if !errors.Is(err, tc.wantErr) {
					t.Fatalf("err=%v, want %v", err, tc.wantErr)
				}
				if selected != nil {
					t.Fatalf("rejected policy reserved candidate %+v", selected)
				}
				if lc := leaseCount(m); lc != 0 {
					t.Fatalf("leaseCount=%d after rejection, want 0", lc)
				}
				for _, id := range []string{recoveryAvoidID, recoveryAltID} {
					if inflight, _ := providerResourceCounts(m, recoveryNodeID, id); inflight != 0 {
						t.Fatalf("provider %s in-flight=%d after rejection, want 0", id, inflight)
					}
				}
				return
			}

			if err != nil {
				t.Fatalf("admit err=%v, want dispatch of %s", err, tc.wantProviderID)
			}
			if selected == nil || selected.providerID != tc.wantProviderID {
				t.Fatalf("selected=%+v, want providerID=%s", selected, tc.wantProviderID)
			}
			if lc := leaseCount(m); lc != 1 {
				t.Fatalf("leaseCount=%d after dispatch, want exactly 1", lc)
			}
			if inflight, _ := providerResourceCounts(m, recoveryNodeID, tc.wantProviderID); inflight != 1 {
				t.Fatalf("provider %s in-flight=%d after dispatch, want 1", tc.wantProviderID, inflight)
			}

			// Release the lease and confirm every counter settles.
			m.releaseLease(selected.leaseID, "test-settle")
			if lc := leaseCount(m); lc != 0 {
				t.Fatalf("leaseCount=%d after release, want 0", lc)
			}
			if inflight, _ := providerResourceCounts(m, recoveryNodeID, tc.wantProviderID); inflight != 0 {
				t.Fatalf("provider %s in-flight=%d after release, want 0", tc.wantProviderID, inflight)
			}
		})
	}
}

// recoveryAdmitResult carries a queued admission outcome back to the test body.
type recoveryAdmitResult struct {
	candidate *candidateNode
	err       error
}

// TestProviderRecoverySelectionQueuedReresolution proves the queued path
// reapplies the identical request-local recovery policy after a runtime-health
// overlay change lands between enqueue and pump: an eligible alternate that
// disappears either promotes the avoided provider under explicit fallback or
// yields a typed terminal rejection when fallback is not permitted.
func TestProviderRecoverySelectionQueuedReresolution(t *testing.T) {
	for _, tc := range []struct {
		name           string
		fallback       bool
		wantProviderID string
		wantErr        error
	}{
		{name: "fallback_true_promotes_avoided", fallback: true, wantProviderID: recoveryAvoidID},
		{name: "fallback_false_rejects", fallback: false, wantErr: ErrProviderPoolCandidateRejected},
	} {
		t.Run(tc.name, func(t *testing.T) {
			m, entry := newRecoveryQueueFixture(t)

			// Occupy the alternate's only slot so a recovery request that prefers
			// it must queue instead of dispatching immediately.
			filler, _, err := m.admitWithRecovery(t.Context(), "filler-group", "", recoveryServed,
				[]candidateNode{recoveryCandidate(entry, recoveryAltID)}, groupPolicy{}, nil, false, true, recoveryCandidatePolicy{})
			if err != nil || filler == nil || filler.providerID != recoveryAltID {
				t.Fatalf("filler admit: candidate=%+v err=%v", filler, err)
			}

			resolver := func() ([]candidateNode, error) {
				return []candidateNode{
					recoveryCandidate(entry, recoveryAvoidID),
					recoveryCandidate(entry, recoveryAltID),
				}, nil
			}

			resultCh := make(chan recoveryAdmitResult, 1)
			ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
			defer cancel()
			go func() {
				candidate, _, admitErr := m.admitWithRecovery(ctx, recoveryGroupKey, "", recoveryServed,
					[]candidateNode{recoveryCandidate(entry, recoveryAvoidID), recoveryCandidate(entry, recoveryAltID)},
					groupPolicy{}, resolver, false, true,
					recoveryCandidatePolicy{avoidProviderID: recoveryAvoidID, allowAvoidedProviderFallback: tc.fallback})
				resultCh <- recoveryAdmitResult{candidate: candidate, err: admitErr}
			}()

			requireProviderPoolPending(t, m, 1)

			// Overlay change before pump: the alternate becomes runtime-unavailable.
			markRecoveryUnavailable(m, entry, recoveryAltID)
			m.mu.Lock()
			m.pumpAllLocked()
			m.mu.Unlock()

			result := <-resultCh
			if tc.wantErr != nil {
				if !errors.Is(result.err, tc.wantErr) {
					t.Fatalf("queued err=%v, want %v", result.err, tc.wantErr)
				}
				if result.candidate != nil {
					t.Fatalf("queued rejection reserved candidate %+v", result.candidate)
				}
				// Only the filler lease remains.
				if lc := leaseCount(m); lc != 1 {
					t.Fatalf("leaseCount=%d after queued rejection, want 1 (filler only)", lc)
				}
			} else {
				if result.err != nil {
					t.Fatalf("queued admit err=%v, want dispatch of %s", result.err, tc.wantProviderID)
				}
				if result.candidate == nil || result.candidate.providerID != tc.wantProviderID {
					t.Fatalf("queued selected=%+v, want providerID=%s", result.candidate, tc.wantProviderID)
				}
				if inflight, _ := providerResourceCounts(m, recoveryNodeID, tc.wantProviderID); inflight != 1 {
					t.Fatalf("provider %s in-flight=%d after queued dispatch, want 1", tc.wantProviderID, inflight)
				}
				if lc := leaseCount(m); lc != 2 {
					t.Fatalf("leaseCount=%d after queued dispatch, want 2 (filler + recovery)", lc)
				}
				m.releaseLease(result.candidate.leaseID, "test-settle")
			}

			// The pending queue must have drained in both branches.
			m.mu.Lock()
			pending := m.pendingProviderPoolCountLocked()
			m.mu.Unlock()
			if pending != 0 {
				t.Fatalf("pending=%d after pump, want 0", pending)
			}

			m.releaseLease(filler.leaseID, "test-cleanup")
			if lc := leaseCount(m); lc != 0 {
				t.Fatalf("leaseCount=%d after full cleanup, want 0", lc)
			}
		})
	}
}

// TestProviderRecoverySelectionServiceDispatchPrefersAlternate exercises the
// full SubmitProviderPool surface over net.Pipe: with the avoided provider
// suppressed, the returned DispatchInfo names the alternate, and exactly one
// provider tunnel request reaches the node (the avoided provider is never
// dispatched).
func TestProviderRecoverySelectionServiceDispatchPrefersAlternate(t *testing.T) {
	edgeConn, nodeConn := net.Pipe()
	t.Cleanup(func() {
		_ = edgeConn.Close()
		_ = nodeConn.Close()
	})

	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.ProviderTunnelRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}
	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

	var capturedMu sync.Mutex
	var capturedCount int
	var capturedAdapter string
	toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
		capturedMu.Lock()
		capturedCount++
		// The alternate and avoided providers use distinct adapter instances
		// (vllm-b vs vllm-a) but the same served target, so the wire adapter is
		// the identity that proves which provider was dispatched.
		capturedAdapter = req.GetAdapter()
		capturedMu.Unlock()
	})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      recoveryNodeID,
		Runtime: config.RuntimeConf{Concurrency: 4},
		Adapters: config.AdaptersConf{
			VllmInstances: []config.VllmInstanceConf{
				{Name: recoveryAvoidAdap, Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				{Name: recoveryAltAdap, Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
			},
		},
		Providers: []config.NodeProviderConf{
			{ID: recoveryAvoidID, Adapter: recoveryAvoidAdap, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
			{ID: recoveryAltID, Adapter: recoveryAltAdap, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
		},
	})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{
		NodeID:                       recoveryNodeID,
		LifecycleState:               edgenode.LifecycleConnected,
		Client:                       edgeClient,
		CredentialRecipientKeyID:     "recipient-recovery",
		CredentialRecipientPublicKey: make([]byte, 32),
	})

	svc := New(reg, edgeevents.NewBus())
	svc.SetNodeStore(store)
	svc.SetModelCatalog([]config.ModelCatalogEntry{
		{ID: recoveryGroupKey, Providers: map[string]string{recoveryAvoidID: recoveryServed, recoveryAltID: recoveryServed}},
	})

	result, err := svc.SubmitProviderPool(context.Background(), ProviderPoolDispatchRequest{
		Run: SubmitRunRequest{
			ModelGroupKey: recoveryGroupKey,
			ProviderPool:  true,
			Background:    true,
		},
		AvoidProviderID: recoveryAvoidID,
	})
	if err != nil {
		t.Fatalf("SubmitProviderPool: %v", err)
	}
	if result == nil || result.Path != ProviderPoolPathTunnel {
		t.Fatalf("result=%+v, want tunnel path", result)
	}
	if result.Tunnel != nil {
		defer result.Tunnel.Close()
	}
	if result.DispatchInfo.ProviderID != recoveryAltID {
		t.Fatalf("DispatchInfo.ProviderID=%q, want %q (avoided provider must be suppressed)", result.DispatchInfo.ProviderID, recoveryAltID)
	}

	waitForCondition(t, func() bool {
		capturedMu.Lock()
		defer capturedMu.Unlock()
		return capturedCount == 1
	}, "expected exactly one provider tunnel request to reach the node")

	capturedMu.Lock()
	defer capturedMu.Unlock()
	if capturedCount != 1 {
		t.Fatalf("captured %d provider tunnel requests, want exactly 1", capturedCount)
	}
	if capturedAdapter != recoveryAltAdap {
		t.Fatalf("wire adapter=%q, want %q (avoided provider was dispatched)", capturedAdapter, recoveryAltAdap)
	}
	if got := inflightRunCount(svc.queue); got != 1 {
		t.Fatalf("inflight run count=%d after single dispatch, want 1", got)
	}
}

// TestProviderRecoverySelectionServiceQueuedReresolution crosses the public
// SubmitProviderPool surface with its default resolver. The alternate first
// fills its capacity; a queued recovery request has no operation or custom
// predicate, then re-resolves against a changed live catalog when the filler
// lease releases. This proves the request-local recovery policy survives the
// public queued path rather than only a queue-core fixture.
func TestProviderRecoverySelectionServiceQueuedReresolution(t *testing.T) {
	for _, tc := range []struct {
		name             string
		fallback         bool
		wantProviderID   string
		wantRecoveryWire int
		wantErr          error
	}{
		{name: "fallback_true_dispatches_the_now_only_avoided_provider", fallback: true, wantProviderID: recoveryAvoidID, wantRecoveryWire: 2},
		{name: "fallback_false_terminates_without_avoided_dispatch", fallback: false, wantRecoveryWire: 1, wantErr: ErrProviderPoolCandidateRejected},
	} {
		t.Run(tc.name, func(t *testing.T) {
			edgeConn, nodeConn := net.Pipe()
			t.Cleanup(func() {
				_ = edgeConn.Close()
				_ = nodeConn.Close()
			})

			parserMap := toki.ParserMap{
				toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.ProviderTunnelRequest{}
					return m, proto.Unmarshal(b, m)
				},
			}
			edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
			nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)

			var capturedMu sync.Mutex
			var capturedAdapters []string
			toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(req *iop.ProviderTunnelRequest) {
				capturedMu.Lock()
				capturedAdapters = append(capturedAdapters, req.GetAdapter())
				capturedMu.Unlock()
			})

			store := edgenode.NewNodeStore()
			store.Add(&edgenode.NodeRecord{
				ID:      recoveryNodeID,
				Runtime: config.RuntimeConf{Concurrency: 4},
				Adapters: config.AdaptersConf{VllmInstances: []config.VllmInstanceConf{
					{Name: recoveryAvoidAdap, Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
					{Name: recoveryAltAdap, Enabled: true, Endpoint: "http://127.0.0.1:8001/v1"},
				}},
				Providers: []config.NodeProviderConf{
					{ID: recoveryAvoidID, Adapter: recoveryAvoidAdap, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
					{ID: recoveryAltID, Adapter: recoveryAltAdap, Type: "vllm", Models: []string{recoveryServed}, Health: "available", Capacity: 1},
				},
			})

			reg := edgenode.NewRegistry()
			reg.Register(&edgenode.NodeEntry{
				NodeID:                       recoveryNodeID,
				LifecycleState:               edgenode.LifecycleConnected,
				Client:                       edgeClient,
				CredentialRecipientKeyID:     "recipient-recovery",
				CredentialRecipientPublicKey: make([]byte, 32),
			})
			svc := New(reg, edgeevents.NewBus())
			svc.SetNodeStore(store)

			// Fill the alternate before the recovery request sees both candidates.
			svc.SetModelCatalog([]config.ModelCatalogEntry{{
				ID: recoveryGroupKey, Providers: map[string]string{recoveryAltID: recoveryServed},
			}})
			filler, err := svc.SubmitProviderPool(t.Context(), ProviderPoolDispatchRequest{
				Run: SubmitRunRequest{ModelGroupKey: recoveryGroupKey, ProviderPool: true, Background: true},
			})
			if err != nil || filler == nil || filler.DispatchInfo.ProviderID != recoveryAltID {
				t.Fatalf("alternate filler: result=%+v err=%v", filler, err)
			}
			waitForCondition(t, func() bool {
				capturedMu.Lock()
				defer capturedMu.Unlock()
				return len(capturedAdapters) == 1
			}, "expected one alternate filler tunnel dispatch")

			// The recovery request has the default empty operation and no custom
			// predicate. While the alternate is capacity-full it must remain queued.
			svc.SetModelCatalog([]config.ModelCatalogEntry{{
				ID: recoveryGroupKey, Providers: map[string]string{recoveryAvoidID: recoveryServed, recoveryAltID: recoveryServed},
			}})
			resultCh := make(chan *ProviderPoolDispatchResult, 1)
			errCh := make(chan error, 1)
			go func() {
				result, submitErr := svc.SubmitProviderPool(t.Context(), ProviderPoolDispatchRequest{
					Run:                          SubmitRunRequest{ModelGroupKey: recoveryGroupKey, ProviderPool: true, Background: true},
					AvoidProviderID:              recoveryAvoidID,
					AllowAvoidedProviderFallback: tc.fallback,
				})
				resultCh <- result
				errCh <- submitErr
			}()
			requireProviderPoolPending(t, svc.queue, 1)

			// Re-resolution must observe the changed catalog, not the enqueue-time
			// slice. Releasing the filler is the production queue pump trigger.
			svc.SetModelCatalog([]config.ModelCatalogEntry{{
				ID: recoveryGroupKey, Providers: map[string]string{recoveryAvoidID: recoveryServed},
			}})
			svc.queue.releaseRun(filler.DispatchInfo.RunID, "test-release-filler")

			result := <-resultCh
			err = <-errCh
			if tc.wantErr != nil {
				if !errors.Is(err, tc.wantErr) {
					t.Fatalf("queued SubmitProviderPool err=%v, want %v", err, tc.wantErr)
				}
				if result != nil {
					t.Fatalf("terminal recovery returned result=%+v", result)
				}
			} else {
				if err != nil || result == nil || result.DispatchInfo.ProviderID != tc.wantProviderID {
					t.Fatalf("queued recovery result=%+v err=%v, want provider %q", result, err, tc.wantProviderID)
				}
				if lc := leaseCount(svc.queue); lc != 1 {
					t.Fatalf("leaseCount=%d after recovery dispatch, want exactly 1", lc)
				}
				svc.queue.releaseRun(result.DispatchInfo.RunID, "test-release-recovery")
			}

			waitForCondition(t, func() bool {
				capturedMu.Lock()
				defer capturedMu.Unlock()
				return len(capturedAdapters) == tc.wantRecoveryWire
			}, "unexpected provider tunnel dispatch count")
			capturedMu.Lock()
			gotAdapters := append([]string(nil), capturedAdapters...)
			capturedMu.Unlock()
			if gotAdapters[0] != recoveryAltAdap {
				t.Fatalf("filler adapter=%q, want %q", gotAdapters[0], recoveryAltAdap)
			}
			if tc.fallback && gotAdapters[1] != recoveryAvoidAdap {
				t.Fatalf("fallback adapter=%q, want permitted avoided adapter %q", gotAdapters[1], recoveryAvoidAdap)
			}
			if lc := leaseCount(svc.queue); lc != 0 {
				t.Fatalf("leaseCount=%d after cleanup, want 0", lc)
			}
		})
	}
}

// TestProviderRecoverySelectionServiceRejectsUnavailableOrUnknownAvoidedProvider
// covers the same-only terminal branches through SubmitProviderPool. Neither a
// runtime-unavailable avoided provider nor a configured-unknown one may reserve
// a lease or emit a provider tunnel request, even when same-provider fallback is
// explicitly permitted.
func TestProviderRecoverySelectionServiceRejectsUnavailableOrUnknownAvoidedProvider(t *testing.T) {
	for _, tc := range []struct {
		name               string
		configuredHealth   string
		markRuntimeOffline bool
		wantErr            error
	}{
		{name: "runtime_unavailable", configuredHealth: "available", markRuntimeOffline: true, wantErr: errProviderUnavailable},
		{name: "configured_unknown", configuredHealth: "unknown"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			edgeConn, nodeConn := net.Pipe()
			t.Cleanup(func() {
				_ = edgeConn.Close()
				_ = nodeConn.Close()
			})
			parserMap := toki.ParserMap{
				toki.TypeNameOf(&iop.ProviderTunnelRequest{}): func(b []byte) (proto.Message, error) {
					m := &iop.ProviderTunnelRequest{}
					return m, proto.Unmarshal(b, m)
				},
			}
			edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
			nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
			var capturedMu sync.Mutex
			captured := 0
			toki.AddListenerTyped[*iop.ProviderTunnelRequest](&nodeClient.Communicator, func(*iop.ProviderTunnelRequest) {
				capturedMu.Lock()
				captured++
				capturedMu.Unlock()
			})

			store := edgenode.NewNodeStore()
			store.Add(&edgenode.NodeRecord{
				ID:      recoveryNodeID,
				Runtime: config.RuntimeConf{Concurrency: 1},
				Adapters: config.AdaptersConf{VllmInstances: []config.VllmInstanceConf{{
					Name: recoveryAvoidAdap, Enabled: true, Endpoint: "http://127.0.0.1:8000/v1",
				}}},
				Providers: []config.NodeProviderConf{{
					ID: recoveryAvoidID, Adapter: recoveryAvoidAdap, Type: "vllm", Models: []string{recoveryServed}, Health: tc.configuredHealth, Capacity: 1,
				}},
			})
			reg := edgenode.NewRegistry()
			entry := &edgenode.NodeEntry{NodeID: recoveryNodeID, LifecycleState: edgenode.LifecycleConnected, Client: edgeClient}
			reg.Register(entry)
			svc := New(reg, edgeevents.NewBus())
			svc.SetNodeStore(store)
			svc.SetModelCatalog([]config.ModelCatalogEntry{{
				ID: recoveryGroupKey, Providers: map[string]string{recoveryAvoidID: recoveryServed},
			}})
			if tc.markRuntimeOffline {
				markRecoveryUnavailable(svc.queue, entry, recoveryAvoidID)
			}

			result, err := svc.SubmitProviderPool(t.Context(), ProviderPoolDispatchRequest{
				Run:                          SubmitRunRequest{ModelGroupKey: recoveryGroupKey, ProviderPool: true, Background: true},
				AvoidProviderID:              recoveryAvoidID,
				AllowAvoidedProviderFallback: true,
			})
			if tc.wantErr != nil {
				if !errors.Is(err, tc.wantErr) {
					t.Fatalf("SubmitProviderPool err=%v, want %v", err, tc.wantErr)
				}
			} else if err == nil {
				t.Fatal("configured-unknown provider unexpectedly dispatched")
			}
			if result != nil {
				t.Fatalf("terminal branch returned result=%+v", result)
			}
			if lc := leaseCount(svc.queue); lc != 0 {
				t.Fatalf("leaseCount=%d after terminal branch, want 0", lc)
			}
			time.Sleep(20 * time.Millisecond)
			capturedMu.Lock()
			defer capturedMu.Unlock()
			if captured != 0 {
				t.Fatalf("captured %d provider tunnel requests after terminal branch, want 0", captured)
			}
		})
	}
}