sub2api/backend/internal/service/account_proxy_selection_test.go

package service

import (
	"context"
	"errors"
	"testing"
	"time"
)

type randomProxySelectorStub struct {
	proxy *Proxy
	err   error
	calls int
}

// randomHydrationRepo only supplies the optional selector capability. The
// embedded interface keeps this test independent of the many account-repository
// methods unrelated to request-time proxy resolution.
type randomHydrationRepo struct {
	AccountRepository
	selector *randomProxySelectorStub
}

func (r *randomHydrationRepo) SelectRandomActiveProxy(ctx context.Context) (*Proxy, error) {
	return r.selector.SelectRandomActiveProxy(ctx)
}

func (s *randomProxySelectorStub) SelectRandomActiveProxy(context.Context) (*Proxy, error) {
	s.calls++
	return s.proxy, s.err
}

// rotatingProxySelectorStub cycles through a fixed pool so tests can assert
// that repeated request-time resolution actually rotates proxies instead of
// sticking to the first one sampled.
type rotatingProxySelectorStub struct {
	proxies []*Proxy
	calls   int
}

func (s *rotatingProxySelectorStub) SelectRandomActiveProxy(context.Context) (*Proxy, error) {
	if len(s.proxies) == 0 {
		return nil, nil
	}
	proxy := s.proxies[s.calls%len(s.proxies)]
	s.calls++
	return proxy, nil
}

func randomProxyAccount() *Account {
	return &Account{
		ID:    42,
		Extra: map[string]any{ProxyModeExtraKey: ProxyModeRandom},
	}
}

func TestResolveRandomProxyAssignsEphemeralProxy(t *testing.T) {
	proxy := &Proxy{ID: 7, Status: StatusActive}
	selector := &randomProxySelectorStub{proxy: proxy}
	account := randomProxyAccount()

	if err := ResolveRandomProxy(context.Background(), account, selector); err != nil {
		t.Fatalf("ResolveRandomProxy() error = %v", err)
	}
	if selector.calls != 1 {
		t.Fatalf("selector calls = %d, want 1", selector.calls)
	}
	if account.Proxy != proxy || account.ProxyID == nil || *account.ProxyID != proxy.ID {
		t.Fatalf("account proxy was not assigned: proxy=%#v proxy_id=%v", account.Proxy, account.ProxyID)
	}
}

func TestResolveRandomProxyFailsClosedWhenPoolUnavailable(t *testing.T) {
	cases := []struct {
		name string
		stub *randomProxySelectorStub
	}{
		{name: "nil selector", stub: nil},
		{name: "empty pool", stub: &randomProxySelectorStub{}},
		{name: "inactive proxy", stub: &randomProxySelectorStub{proxy: &Proxy{ID: 1, Status: StatusDisabled}}},
		{name: "expired proxy", stub: &randomProxySelectorStub{proxy: &Proxy{ID: 1, Status: StatusActive, ExpiresAt: accountProxyTimePtr(time.Now().Add(-time.Minute))}}},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			account := randomProxyAccount()
			account.ProxyID = accountProxyInt64Ptr(99)
			account.Proxy = &Proxy{ID: 99, Status: StatusActive}
			var err error
			if tc.stub == nil {
				err = ResolveRandomProxy(context.Background(), account, nil)
			} else {
				err = ResolveRandomProxy(context.Background(), account, tc.stub)
			}
			if !errors.Is(err, ErrRandomProxyUnavailable) {
				t.Fatalf("error = %v, want ErrRandomProxyUnavailable", err)
			}
			// A failed refresh must not leave a stale fixed proxy usable.
			if account.ProxyID != nil || account.Proxy != nil {
				t.Fatalf("failed selection left a stale proxy association")
			}
		})
	}
}

func TestResolveRandomProxyRotatesAcrossCalls(t *testing.T) {
	proxies := []*Proxy{
		{ID: 21, Status: StatusActive},
		{ID: 22, Status: StatusActive},
		{ID: 23, Status: StatusActive},
	}
	selector := &rotatingProxySelectorStub{proxies: proxies}

	seen := map[int64]int{}
	const resolves = 30
	for i := 0; i < resolves; i++ {
		account := randomProxyAccount()
		if err := ResolveRandomProxy(context.Background(), account, selector); err != nil {
			t.Fatalf("ResolveRandomProxy() error = %v", err)
		}
		if account.Proxy == nil || account.ProxyID == nil {
			t.Fatal("rotating selection left the account without a proxy")
		}
		seen[*account.ProxyID]++
	}
	for _, proxy := range proxies {
		if seen[proxy.ID] == 0 {
			t.Fatalf("proxy %d was never selected in %d resolves (seen=%v)", proxy.ID, resolves, seen)
		}
	}
}

func TestNormalizeProxyModeExtra(t *testing.T) {
	cases := []struct {
		name  string
		input map[string]any
		want  map[string]any
	}{
		{name: "nil map stays nil", input: nil, want: nil},
		{name: "missing key untouched", input: map[string]any{"other": 1}, want: map[string]any{"other": 1}},
		{name: "random kept", input: map[string]any{ProxyModeExtraKey: "random"}, want: map[string]any{ProxyModeExtraKey: "random"}},
		{name: "case and whitespace canonicalized", input: map[string]any{ProxyModeExtraKey: " Random "}, want: map[string]any{ProxyModeExtraKey: "random"}},
		{name: "other string dropped", input: map[string]any{ProxyModeExtraKey: "fixed", "other": 1}, want: map[string]any{"other": 1}},
		{name: "empty string dropped", input: map[string]any{ProxyModeExtraKey: ""}, want: map[string]any{}},
		{name: "non-string dropped", input: map[string]any{ProxyModeExtraKey: 7}, want: map[string]any{}},
		{name: "nil value dropped", input: map[string]any{ProxyModeExtraKey: nil}, want: map[string]any{}},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got := NormalizeProxyModeExtra(tc.input)
			if len(got) != len(tc.want) {
				t.Fatalf("NormalizeProxyModeExtra() = %v, want %v", got, tc.want)
			}
			for key, wantValue := range tc.want {
				if got[key] != wantValue {
					t.Fatalf("NormalizeProxyModeExtra()[%q] = %v, want %v", key, got[key], wantValue)
				}
			}
		})
	}
}

func TestResolveRandomProxyIgnoresNonRandomAccount(t *testing.T) {
	selector := &randomProxySelectorStub{proxy: &Proxy{ID: 1, Status: StatusActive}}
	account := &Account{ID: 1, Extra: map[string]any{ProxyModeExtraKey: "fixed"}}

	if err := ResolveRandomProxy(context.Background(), account, selector); err != nil {
		t.Fatalf("ResolveRandomProxy() error = %v", err)
	}
	if selector.calls != 0 {
		t.Fatalf("selector calls = %d, want 0", selector.calls)
	}
}

func TestGatewayHydrationResolvesRandomProxyWithoutSnapshot(t *testing.T) {
	selector := &randomProxySelectorStub{proxy: &Proxy{ID: 11, Status: StatusActive}}
	repo := &randomHydrationRepo{selector: selector}
	service := &GatewayService{accountRepo: repo}
	account := randomProxyAccount()
	account.ProxyID = accountProxyInt64Ptr(99)
	account.Proxy = &Proxy{ID: 99, Status: StatusActive}

	hydrated, err := service.hydrateSelectedAccount(context.Background(), account)
	if err != nil {
		t.Fatalf("hydrateSelectedAccount() error = %v", err)
	}
	if hydrated != account || hydrated.Proxy == nil || hydrated.Proxy.ID != 11 {
		t.Fatalf("hydrated proxy = %#v, want proxy 11", hydrated.Proxy)
	}
}

func TestOpenAIHydrationResolvesRandomProxyWithoutSnapshot(t *testing.T) {
	selector := &randomProxySelectorStub{proxy: &Proxy{ID: 12, Status: StatusActive}}
	repo := &randomHydrationRepo{selector: selector}
	service := &OpenAIGatewayService{accountRepo: repo}

	hydrated, err := service.hydrateSelectedAccount(context.Background(), randomProxyAccount())
	if err != nil {
		t.Fatalf("hydrateSelectedAccount() error = %v", err)
	}
	if hydrated == nil || hydrated.Proxy == nil || hydrated.Proxy.ID != 12 {
		t.Fatalf("hydrated proxy = %#v, want proxy 12", hydrated.Proxy)
	}
}

func TestGeminiHydrationResolvesRandomProxyWithoutSnapshot(t *testing.T) {
	selector := &randomProxySelectorStub{proxy: &Proxy{ID: 13, Status: StatusActive}}
	repo := &randomHydrationRepo{selector: selector}
	service := &GeminiMessagesCompatService{accountRepo: repo}

	hydrated, err := service.hydrateSelectedAccount(context.Background(), randomProxyAccount())
	if err != nil {
		t.Fatalf("hydrateSelectedAccount() error = %v", err)
	}
	if hydrated == nil || hydrated.Proxy == nil || hydrated.Proxy.ID != 13 {
		t.Fatalf("hydrated proxy = %#v, want proxy 13", hydrated.Proxy)
	}
}

func accountProxyTimePtr(value time.Time) *time.Time { return &value }

func accountProxyInt64Ptr(value int64) *int64 { return &value }