ollama-proxy/internal/proxy/balancer_test.go

package proxy

import (
	"sync"
	"testing"
	"time"

	"github.com/Atte149/ollama-proxy/internal/config"
)

func mkAccts(n int) []config.Account {
	out := make([]config.Account, n)
	for i := range out {
		out[i] = config.Account{ID: "id" + string(rune('a'+i)), Name: "a" + string(rune('a'+i)), APIKey: "k" + string(rune('a'+i))}
	}
	return out
}

func TestBalancer_Next_RoundRobin(t *testing.T) {
	b := NewBalancer(mkAccts(3), 60*time.Second)
	seen := make(map[string]int)
	for i := 0; i < 9; i++ {
		a, err := b.Next()
		if err != nil {
			t.Fatalf("Next %d: %v", i, err)
		}
		seen[a.ID]++
	}
	// 9 calls across 3 accounts should hit each exactly 3 times (round-robin).
	for id, n := range seen {
		if n != 3 {
			t.Errorf("account %s got %d hits, want 3", id, n)
		}
	}
}

func TestBalancer_Next_NoAccounts(t *testing.T) {
	b := NewBalancer(nil, 60*time.Second)
	_, err := b.Next()
	if _, ok := err.(ErrNoAccounts); !ok {
		t.Errorf("Next on empty balancer: %v, want ErrNoAccounts", err)
	}
}

func TestBalancer_Next_SkipsCooldown(t *testing.T) {
	accts := mkAccts(3)
	b := NewBalancer(accts, 60*time.Second)
	// Put the first account in cooldown; the next 3 calls should never pick it.
	b.MarkCooldown(accts[0].ID)
	for i := 0; i < 3; i++ {
		a, err := b.Next()
		if err != nil {
			t.Fatalf("Next %d: %v", i, err)
		}
		if a.ID == accts[0].ID {
			t.Errorf("Next %d returned cooldown account %s", i, a.ID)
		}
	}
}

func TestBalancer_Next_AllCooldown(t *testing.T) {
	accts := mkAccts(2)
	b := NewBalancer(accts, 60*time.Second)
	b.MarkCooldown(accts[0].ID)
	b.MarkCooldown(accts[1].ID)
	_, err := b.Next()
	if _, ok := err.(ErrAllCooldown); !ok {
		t.Errorf("Next with all in cooldown: %v, want ErrAllCooldown", err)
	}
}

func TestBalancer_ClearCooldown(t *testing.T) {
	accts := mkAccts(2)
	b := NewBalancer(accts, 60*time.Second)
	b.MarkCooldown(accts[0].ID)
	b.ClearCooldown(accts[0].ID)
	// After clearing, the first account must be selectable again. Spin until
	// we see it — round-robin visits each account within len(accts) calls.
	seen := false
	for i := 0; i < len(accts)*3; i++ {
		a, err := b.Next()
		if err != nil {
			t.Fatalf("Next %d: %v", i, err)
		}
		if a.ID == accts[0].ID {
			seen = true
			break
		}
	}
	if !seen {
		t.Error("cleared account was never returned by Next")
	}
}

func TestBalancer_Next_Concurrent(t *testing.T) {
	b := NewBalancer(mkAccts(3), 60*time.Second)
	var wg sync.WaitGroup
	const goroutines = 16
	const perG = 100
	results := make(chan string, goroutines*perG)
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < perG; i++ {
				a, err := b.Next()
				if err != nil {
					t.Errorf("Next: %v", err)
					return
				}
				results <- a.ID
			}
		}()
	}
	wg.Wait()
	close(results)
	// Sanity: every result is one of the known account IDs. We don't assert
	// strict fairness under contention (atomic counter gives best-effort RR).
	known := map[string]bool{"ida": true, "idb": true, "idc": true}
	count := 0
	for id := range results {
		if !known[id] {
			t.Errorf("unknown account id returned: %s", id)
		}
		count++
	}
	if count != goroutines*perG {
		t.Errorf("got %d results, want %d", count, goroutines*perG)
	}
}

func TestBalancer_Status(t *testing.T) {
	accts := mkAccts(2)
	b := NewBalancer(accts, 60*time.Second)
	b.MarkCooldown(accts[0].ID)
	st := b.Status()
	if !st[accts[0].ID].InCooldown {
		t.Error("account 0 should be in cooldown")
	}
	if st[accts[1].ID].InCooldown {
		t.Error("account 1 should not be in cooldown")
	}
	if st[accts[0].ID].Last429.IsZero() {
		t.Error("account 0 should have Last429 set")
	}
}

func TestBalancer_CooldownExpires(t *testing.T) {
	accts := mkAccts(1)
	b := NewBalancer(accts, 50*time.Millisecond)
	b.MarkCooldown(accts[0].ID)
	// Immediately: account is in cooldown.
	if _, err := b.Next(); err == nil {
		t.Fatal("Next succeeded immediately after MarkCooldown")
	}
	time.Sleep(60 * time.Millisecond)
	if _, err := b.Next(); err != nil {
		t.Errorf("Next after cooldown expiry: %v", err)
	}
}