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)),
			Provider: config.ProviderOllamaCloud,
		}
	}
	return out
}

func mkMixedAccts() []config.Account {
	return []config.Account{
		{ID: "idO1", Name: "o1", APIKey: "k1", Provider: config.ProviderOllamaCloud},
		{ID: "idO2", Name: "o2", APIKey: "k2", Provider: config.ProviderOllamaCloud},
		{ID: "idG1", Name: "g1", APIKey: "k3", Provider: config.ProviderOpenCodeGo},
		{ID: "idG2", Name: "g2", APIKey: "k4", Provider: config.ProviderOpenCodeGo},
	}
}

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)
	}
}

func TestBalancer_NextFor_FiltersByProvider(t *testing.T) {
	accts := mkMixedAccts()
	b := NewBalancer(accts, 60*time.Second)
	// Request only OpenCode Go accounts: every result must be a Go account.
	seen := map[string]int{}
	for i := 0; i < 8; i++ {
		a, err := b.NextFor([]config.ProviderType{config.ProviderOpenCodeGo})
		if err != nil {
			t.Fatalf("NextFor %d: %v", i, err)
		}
		if a.Provider != config.ProviderOpenCodeGo {
			t.Errorf("NextFor returned provider %q, want opencode-go", a.Provider)
		}
		seen[a.ID]++
	}
	if len(seen) != 2 {
		t.Errorf("expected 2 distinct go accounts, got %d", len(seen))
	}
}

func TestBalancer_NextFor_MixedPool(t *testing.T) {
	accts := mkMixedAccts()
	b := NewBalancer(accts, 60*time.Second)
	// Eligible = both providers: all 4 accounts should be reachable.
	seen := map[string]bool{}
	for i := 0; i < 16; i++ {
		a, err := b.NextFor([]config.ProviderType{
			config.ProviderOllamaCloud, config.ProviderOpenCodeGo,
		})
		if err != nil {
			t.Fatalf("NextFor %d: %v", i, err)
		}
		seen[a.ID] = true
	}
	if len(seen) != 4 {
		t.Errorf("expected 4 distinct accounts in mixed pool, got %d", len(seen))
	}
}

func TestBalancer_NextFor_NilEligibleReturnsAll(t *testing.T) {
	accts := mkMixedAccts()
	b := NewBalancer(accts, 60*time.Second)
	seen := map[config.ProviderType]bool{}
	for i := 0; i < 16; i++ {
		a, err := b.NextFor(nil)
		if err != nil {
			t.Fatalf("NextFor nil %d: %v", i, err)
		}
		seen[a.Provider] = true
	}
	if !seen[config.ProviderOllamaCloud] || !seen[config.ProviderOpenCodeGo] {
		t.Errorf("nil eligible should include both providers, got %v", seen)
	}
}

func TestBalancer_NextFor_NoEligibleAccounts(t *testing.T) {
	// Only Ollama accounts configured, but requesting Go.
	b := NewBalancer(mkAccts(2), 60*time.Second)
	_, err := b.NextFor([]config.ProviderType{config.ProviderOpenCodeGo})
	if _, ok := err.(ErrAllCooldown); !ok && err == nil {
		t.Errorf("NextFor with no eligible accounts: %v, want error or ErrAllCooldown", err)
	}
}