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

package service

import (
	"errors"
	"sync"
	"testing"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

// TestRefreshProviderCapacityAffectsNextDispatch verifies that a provider
// capacity change applied via SetRuntimeConfig is reflected in the candidate
// capacity used for the next dispatch admission, since candidates are rebuilt
// from the live runtime snapshot on each request.
func TestRefreshProviderCapacityAffectsNextDispatch(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"})
	svc := New(reg, nil)

	newStore := func(capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cap",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:       "prov-a",
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{"served-a"},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	catalog := []config.ModelCatalogEntry{{
		ID:        "qwen3.6:35b",
		Providers: map[string]string{"prov-a": "served-a"},
	}}

	req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true}

	svc.SetRuntimeConfig(newStore(2), catalog)
	store, cat := svc.runtimeConfigSnapshot()
	cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve before refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 2 {
		t.Fatalf("before refresh: got candidates=%+v, want one with capacity=2", cands)
	}

	// Refresh raises capacity to 8; the next candidate build must observe it.
	svc.SetRuntimeConfig(newStore(8), catalog)
	store, cat = svc.runtimeConfigSnapshot()
	cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve after refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 8 {
		t.Fatalf("after refresh: got candidates=%+v, want one with capacity=8", cands)
	}
}

// TestRefreshProviderLongCapacityAffectsCandidate verifies that a provider
// long_context_capacity change applied via SetRuntimeConfig is reflected in the
// candidate longContextCapacity used for the next long-context admission.
func TestRefreshProviderLongCapacityAffectsCandidate(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"})
	svc := New(reg, nil)

	newStore := func(longCapacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cap",
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:                  "prov-a",
					Type:                "vllm",
					Category:            config.CategoryAPI,
					Adapter:             "vllm-gpu",
					Models:              []string{"served-a"},
					Health:              "available",
					Capacity:            4,
					TotalContextTokens:  524288,
					LongContextCapacity: longCapacity,
				},
			},
		})
		return store
	}
	catalog := []config.ModelCatalogEntry{{
		ID:                  "qwen3.6:35b",
		ContextWindowTokens: 262144,
		Providers:           map[string]string{"prov-a": "served-a"},
	}}

	req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true, ContextClass: "long"}

	svc.SetRuntimeConfig(newStore(1), catalog)
	store, cat := svc.runtimeConfigSnapshot()
	cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve before refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].longContextCapacity != 1 {
		t.Fatalf("before refresh: got candidates=%+v, want one with longContextCapacity=1", cands)
	}

	// Refresh raises long capacity to 2; the next candidate build must observe it.
	svc.SetRuntimeConfig(newStore(2), catalog)
	store, cat = svc.runtimeConfigSnapshot()
	cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve after refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].longContextCapacity != 2 {
		t.Fatalf("after refresh: got candidates=%+v, want one with longContextCapacity=2", cands)
	}
}

func TestRuntimeConfigSnapshotConcurrentReplace(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-refresh"})
	svc := New(reg, nil)

	newStore := func(providerID, servedModel string, capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID:      "node-refresh",
			Runtime: config.RuntimeConf{Concurrency: capacity},
			Adapters: config.AdaptersConf{
				VllmInstances: []config.VllmInstanceConf{
					{Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"},
				},
			},
			Providers: []config.NodeProviderConf{
				{
					ID:       providerID,
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{servedModel},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	newCatalog := func(providerID, servedModel string) []config.ModelCatalogEntry {
		return []config.ModelCatalogEntry{
			{
				ID: "qwen3.6:35b",
				Providers: map[string]string{
					providerID: servedModel,
				},
			},
		}
	}

	storeA := newStore("prov-a", "served-a", 2)
	catalogA := newCatalog("prov-a", "served-a")
	storeB := newStore("prov-b", "served-b", 4)
	catalogB := newCatalog("prov-b", "served-b")
	svc.SetRuntimeConfig(storeA, catalogA)

	const readers = 8
	const iterations = 200
	start := make(chan struct{})
	errCh := make(chan error, readers)
	var wg sync.WaitGroup

	for i := 0; i < readers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			for j := 0; j < iterations; j++ {
				store, catalog := svc.runtimeConfigSnapshot()
				candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{
					ModelGroupKey: "qwen3.6:35b",
					ProviderPool:  true,
				}, store, catalog)
				if err != nil {
					errCh <- err
					return
				}
				if len(candidates) != 1 {
					errCh <- errors.New("expected exactly one provider candidate")
					return
				}
				_ = svc.ListNodeSnapshots()
			}
		}()
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < iterations; i++ {
			svc.SetRuntimeConfig(storeA, catalogA)
			svc.SetRuntimeConfig(storeB, catalogB)
		}
	}()

	close(start)
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("runtime config snapshot reader observed inconsistent state: %v", err)
		}
	}
}