153 lines
4 KiB
Go
153 lines
4 KiB
Go
|
|
package quota
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"testing"
|
|||
|
|
|
|||
|
|
"forge.secuaas.ovh/olivier/claude-failover/internal/config"
|
|||
|
|
"forge.secuaas.ovh/olivier/claude-failover/internal/state"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// mockTmux for quota tests.
|
|||
|
|
type mockTmux struct {
|
|||
|
|
sessions map[string]bool
|
|||
|
|
paneOutput map[string]string
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func newMockTmux() *mockTmux {
|
|||
|
|
return &mockTmux{
|
|||
|
|
sessions: make(map[string]bool),
|
|||
|
|
paneOutput: make(map[string]string),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (m *mockTmux) HasSession(name string) bool { return m.sessions[name] }
|
|||
|
|
func (m *mockTmux) CreateSession(name, _ string) error { m.sessions[name] = true; return nil }
|
|||
|
|
func (m *mockTmux) KillSession(_ string) error { return nil }
|
|||
|
|
func (m *mockTmux) SendKeys(_, _ string) error { return nil }
|
|||
|
|
func (m *mockTmux) CapturePaneTail(session string, _ int) (string, error) {
|
|||
|
|
return m.paneOutput[session], nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestIsQuotaExhausted verifies pattern matching on pane output.
|
|||
|
|
func TestIsQuotaExhausted(t *testing.T) {
|
|||
|
|
cases := []struct {
|
|||
|
|
input string
|
|||
|
|
want bool
|
|||
|
|
}{
|
|||
|
|
{"You've hit your limit for Claude Pro.", true},
|
|||
|
|
{"rate limit exceeded", true},
|
|||
|
|
{"quota exceeded for this period", true},
|
|||
|
|
{"Usage limit reached", true},
|
|||
|
|
{"Too many requests", true},
|
|||
|
|
{"Some normal output ❯", false},
|
|||
|
|
{"❯ ", false},
|
|||
|
|
{"still running 5s · ", false},
|
|||
|
|
}
|
|||
|
|
for _, c := range cases {
|
|||
|
|
if got := isQuotaExhausted(c.input); got != c.want {
|
|||
|
|
t.Errorf("isQuotaExhausted(%q) = %v, want %v", c.input, got, c.want)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestExtractResetTime parses various reset time formats.
|
|||
|
|
func TestExtractResetTime(t *testing.T) {
|
|||
|
|
cases := []struct {
|
|||
|
|
input string
|
|||
|
|
want string
|
|||
|
|
}{
|
|||
|
|
{"Usage resets 8pm", "8pm"},
|
|||
|
|
{"Your quota resets at 11:30pm", "11:30pm"},
|
|||
|
|
{"resets in 45 minutes", "in 45 minutes"},
|
|||
|
|
{"resets in 2 hours", "in 2 hours"},
|
|||
|
|
{"no reset info here", ""},
|
|||
|
|
}
|
|||
|
|
for _, c := range cases {
|
|||
|
|
if got := extractResetTime(c.input); got != c.want {
|
|||
|
|
t.Errorf("extractResetTime(%q) = %q, want %q", c.input, got, c.want)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestPollTriggersSwitchOnTwoBlockedPool verifies swap trigger for >=2 blocked pool sessions.
|
|||
|
|
func TestPollTriggersSwitchOnTwoBlockedPool(t *testing.T) {
|
|||
|
|
tc := newMockTmux()
|
|||
|
|
tc.sessions["ccl-auto-0"] = true
|
|||
|
|
tc.sessions["ccl-auto-1"] = true
|
|||
|
|
tc.paneOutput["ccl-auto-0"] = "You've hit your limit for Claude Pro."
|
|||
|
|
tc.paneOutput["ccl-auto-1"] = "rate limit exceeded"
|
|||
|
|
|
|||
|
|
s := state.New("")
|
|||
|
|
s.SetActiveAccount("compte1")
|
|||
|
|
|
|||
|
|
cfg := &config.Config{
|
|||
|
|
Pool: config.PoolConfig{
|
|||
|
|
Autonomous: config.AutonomousConfig{Prefix: "ccl-auto-", Max: 2},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
m := New(tc, s, cfg)
|
|||
|
|
m.poll()
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case req := <-m.switchCh:
|
|||
|
|
if req.From != "compte1" {
|
|||
|
|
t.Errorf("expected From=compte1, got %q", req.From)
|
|||
|
|
}
|
|||
|
|
default:
|
|||
|
|
t.Fatal("expected SwitchRequest on channel")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestPollTriggersSwitchOnOneBlockedInteractive verifies swap trigger for >=1 dedicated session.
|
|||
|
|
func TestPollTriggersSwitchOnOneBlockedInteractive(t *testing.T) {
|
|||
|
|
tc := newMockTmux()
|
|||
|
|
tc.sessions["my-session"] = true
|
|||
|
|
tc.paneOutput["my-session"] = "quota exceeded"
|
|||
|
|
|
|||
|
|
s := state.New("")
|
|||
|
|
s.SetActiveAccount("compte1")
|
|||
|
|
|
|||
|
|
cfg := &config.Config{
|
|||
|
|
Pool: config.PoolConfig{
|
|||
|
|
Dedicated: []config.DedicatedSession{{Name: "my-session"}},
|
|||
|
|
Autonomous: config.AutonomousConfig{Max: 0},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
m := New(tc, s, cfg)
|
|||
|
|
m.poll()
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case req := <-m.switchCh:
|
|||
|
|
if req.From != "compte1" {
|
|||
|
|
t.Errorf("expected From=compte1, got %q", req.From)
|
|||
|
|
}
|
|||
|
|
default:
|
|||
|
|
t.Fatal("expected SwitchRequest on channel")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestPollNoTriggerWhenBelowThreshold verifies no swap for a single blocked pool session.
|
|||
|
|
func TestPollNoTriggerWhenBelowThreshold(t *testing.T) {
|
|||
|
|
tc := newMockTmux()
|
|||
|
|
tc.sessions["ccl-auto-0"] = true
|
|||
|
|
tc.sessions["ccl-auto-1"] = true
|
|||
|
|
tc.paneOutput["ccl-auto-0"] = "rate limit exceeded"
|
|||
|
|
tc.paneOutput["ccl-auto-1"] = "❯ " // fine
|
|||
|
|
|
|||
|
|
s := state.New("")
|
|||
|
|
cfg := &config.Config{
|
|||
|
|
Pool: config.PoolConfig{
|
|||
|
|
Autonomous: config.AutonomousConfig{Prefix: "ccl-auto-", Max: 2},
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
m := New(tc, s, cfg)
|
|||
|
|
m.poll()
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case req := <-m.switchCh:
|
|||
|
|
t.Errorf("unexpected SwitchRequest: %+v", req)
|
|||
|
|
default:
|
|||
|
|
// Correct: only 1 blocked pool session, threshold is 2.
|
|||
|
|
}
|
|||
|
|
}
|