llmgw/internal/router/endpoint/endpoint_test.go

package endpoint

import (
	"sync"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
)

type clock struct {
	mu sync.Mutex
	t  time.Time
}

func newClock() *clock {
	return &clock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
}
func (c *clock) now() time.Time {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.t
}
func (c *clock) advance(d time.Duration) {
	c.mu.Lock()
	c.t = c.t.Add(d)
	c.mu.Unlock()
}

func newReg(t *testing.T) (*Registry, *clock) {
	t.Helper()
	r := NewRegistry(Config{StaleAfter: 10 * time.Second, Breaker: breaker.DefaultConfig()})
	c := newClock()
	r.SetClock(c.now)
	return r, c
}

func TestReplaceKeepsLiveState(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", URL: "http://a", Models: []string{"m"}}})
	r.Observe("a", Stats{Running: 5})

	// Re-supplying the same endpoint must not lose its stats.
	r.Replace([]Info{
		{ID: "a", URL: "http://a", Models: []string{"m"}, Weight: 2},
		{ID: "b", URL: "http://b", Models: []string{"m"}},
	})

	snaps := r.Snapshots("m")
	if len(snaps) != 2 {
		t.Fatalf("expected 2 endpoints, got %d", len(snaps))
	}
	if snaps[0].Stats.Running != 5 {
		t.Error("Replace reset the live state")
	}
	if snaps[0].Weight != 2 {
		t.Error("Replace did not update the static info")
	}
}

func TestReplaceRemovesMissingEndpoints(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}, {ID: "b", Models: []string{"m"}}})
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	if got := r.IDs(); len(got) != 1 || got[0] != "a" {
		t.Errorf("endpoint dropped from the list was not removed: %v", got)
	}
}

func TestFilterByModel(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{
		{ID: "a", Models: []string{"llama"}},
		{ID: "b", Models: []string{"qwen"}},
		{ID: "c", Models: []string{"llama", "qwen"}},
	})

	if got := len(r.Snapshots("llama")); got != 2 {
		t.Errorf("got %d endpoints for llama", got)
	}
	if got := len(r.Snapshots("")); got != 3 {
		t.Errorf("empty model should return all, got %d", got)
	}
}

func TestStalenessDetection(t *testing.T) {
	r, c := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	// With no scrape at all the endpoint counts as stale.
	if !r.Snapshots("m")[0].Stale {
		t.Error("not marked stale despite having no measurement")
	}

	r.Observe("a", Stats{Running: 1})
	if r.Snapshots("m")[0].Stale {
		t.Error("a fresh measurement was marked stale")
	}

	c.advance(11 * time.Second)
	s := r.Snapshots("m")[0]
	if !s.Stale {
		t.Error("not marked stale after StaleAfter elapsed")
	}
	if s.Age < 11*time.Second {
		t.Errorf("Age is wrong: %v", s.Age)
	}
}

func TestDrainRateCalculation(t *testing.T) {
	r, c := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	r.Observe("a", Stats{SuccessTotal: 100})
	c.advance(time.Second)
	r.Observe("a", Stats{SuccessTotal: 110}) // 10 requests/sec

	if got := r.Snapshots("m")[0].DrainRate; got <= 0 {
		t.Fatalf("drain rate was not computed: %v", got)
	}

	// Exponential smoothing: one sample approaches the raw rate but stays below it.
	if got := r.Snapshots("m")[0].DrainRate; got > 10 {
		t.Errorf("drain rate exceeded the raw value: %v", got)
	}
}

func TestDrainRateResetsWhenCounterGoesBackwards(t *testing.T) {
	// Counters reset when an engine restarts; that must not yield a negative rate.
	r, c := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	r.Observe("a", Stats{SuccessTotal: 1000})
	c.advance(time.Second)
	r.Observe("a", Stats{SuccessTotal: 1100})
	c.advance(time.Second)
	r.Observe("a", Stats{SuccessTotal: 5}) // restarted

	if got := r.Snapshots("m")[0].DrainRate; got != 0 {
		t.Errorf("drain rate was not reset after the counter went backwards: %v", got)
	}
}

func TestAcquireCountsInflight(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	rel1, ok := r.Acquire("a")
	if !ok {
		t.Fatal("Acquire was refused")
	}
	rel2, ok := r.Acquire("a")
	if !ok {
		t.Fatal("second Acquire was refused")
	}
	if got := r.Inflight("a"); got != 2 {
		t.Errorf("inflight=%d, expected 2", got)
	}

	rel1(true)
	rel2(true)
	if got := r.Inflight("a"); got != 0 {
		t.Errorf("inflight=%d after release", got)
	}
}

func TestReleaseIsIdempotent(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	rel, _ := r.Acquire("a")
	rel(true)
	rel(true) // must be idempotent; the counter must not go negative

	if got := r.Inflight("a"); got != 0 {
		t.Errorf("double release corrupted the counter: %d", got)
	}
}

func TestFailuresOpenTheCircuit(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	for range 3 { // DefaultConfig threshold is 3
		rel, ok := r.Acquire("a")
		if !ok {
			t.Fatal("circuit opened too early")
		}
		rel(false)
	}

	if _, ok := r.Acquire("a"); ok {
		t.Error("circuit did not open after consecutive failures")
	}
	if r.Snapshots("m")[0].Breaker != breaker.Open {
		t.Error("snapshot does not reflect the breaker state")
	}
}

func TestScrapeFailureDoesNotTripBreaker(t *testing.T) {
	// A scrape failure is not a request failure: the endpoint may be serving
	// fine while only its metrics endpoint is unreachable.
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}})

	for range 10 {
		r.ScrapeFailed("a")
	}
	if _, ok := r.Acquire("a"); !ok {
		t.Error("scrape failures opened the circuit")
	}
}

func TestUnknownEndpoint(t *testing.T) {
	r, _ := newReg(t)
	if _, ok := r.Acquire("missing"); ok {
		t.Error("an unregistered endpoint was allowed")
	}
	r.Observe("missing", Stats{}) // must not panic
}

func TestConcurrentAccess(t *testing.T) {
	r, _ := newReg(t)
	r.Replace([]Info{{ID: "a", Models: []string{"m"}}, {ID: "b", Models: []string{"m"}}})

	var wg sync.WaitGroup
	for i := range 100 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			id := "a"
			if i%2 == 0 {
				id = "b"
			}
			if rel, ok := r.Acquire(id); ok {
				rel(i%3 != 0)
			}
			r.Observe(id, Stats{Running: i})
			_ = r.Snapshots("m")
			_ = r.Inflight(id)
		}()
	}
	wg.Wait()
}

func TestLastUsedIsRecordedOnDispatch(t *testing.T) {
	// The fact nothing recorded. Whether a replica is earning the cards it
	// holds is the first question on a fleet with more models than GPUs.
	r := NewRegistry(Config{})
	clock := time.Date(2026, 9, 17, 12, 0, 0, 0, time.UTC)
	r.SetClock(func() time.Time { return clock })
	r.Replace([]Info{{ID: "a", URL: "http://x", Models: []string{"m"}}})

	if len(r.LastUsed()) != 0 {
		t.Fatal("an endpoint nothing has used reported a last-use time")
	}

	release, ok := r.Acquire("a")
	if !ok {
		t.Fatal("acquire refused")
	}
	release(true)

	used := r.LastUsed()
	if got, ok := used["a"]; !ok || !got.Equal(clock) {
		t.Errorf("last used = %v (present=%v), want %v", got, ok, clock)
	}
}

func TestAnUnusedEndpointIsAbsentRatherThanZero(t *testing.T) {
	// A zero time would be the oldest idle replica in any comparison, so a
	// router restart would nominate the whole fleet for eviction at once.
	r := NewRegistry(Config{})
	r.Replace([]Info{
		{ID: "used", URL: "http://x", Models: []string{"m"}},
		{ID: "never", URL: "http://y", Models: []string{"m"}},
	})
	release, _ := r.Acquire("used")
	release(true)

	used := r.LastUsed()
	if _, ok := used["never"]; ok {
		t.Error("an endpoint nothing has used appeared in the map")
	}
	if _, ok := used["used"]; !ok {
		t.Error("the used endpoint is missing")
	}
}

func TestLastUsedSurvivesAnEndpointListRefresh(t *testing.T) {
	// The list is republished whenever anything changes, which on a large
	// fleet is often. Losing usage on every republish would make idleness
	// unmeasurable exactly where it matters most.
	r := NewRegistry(Config{})
	r.Replace([]Info{{ID: "a", URL: "http://x", Models: []string{"m"}}})
	release, _ := r.Acquire("a")
	release(true)
	before := r.LastUsed()["a"]

	// Same endpoint, republished alongside a new one.
	r.Replace([]Info{
		{ID: "a", URL: "http://x", Models: []string{"m"}},
		{ID: "b", URL: "http://y", Models: []string{"m"}},
	})
	if got, ok := r.LastUsed()["a"]; !ok || !got.Equal(before) {
		t.Errorf("last used = %v (present=%v) after a republish, want %v", got, ok, before)
	}
}