internal/auth/session_probe_test.go

package auth

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"
	"testing"
	"time"

	"ds2api/internal/account"
	"ds2api/internal/config"
)

func newProbeTestResolver(t *testing.T, login LoginFunc) (*Resolver, *config.Store) {
	t.Helper()
	t.Setenv("DS2API_CONFIG_JSON", `{
		"keys":["managed-key"],
		"accounts":[{"email":"acc@example.com","password":"pwd"}]
	}`)
	store := config.LoadStore()
	if err := store.UpdateAccountToken("acc@example.com", "session-token"); err != nil {
		t.Fatalf("seed token failed: %v", err)
	}
	return NewResolver(store, account.NewPool(store), login), store
}

func newProbeTestAuth(token string) *RequestAuth {
	return &RequestAuth{
		UseConfigToken: true,
		AccountID:      "acc@example.com",
		Account:        config.Account{Email: "acc@example.com", Password: "pwd", Token: token},
		DeepSeekToken:  token,
		TriedAccounts:  map[string]bool{},
	}
}

func countingLogin(loginCount *int32, token string) LoginFunc {
	return func(_ context.Context, _ config.Account) (string, error) {
		atomic.AddInt32(loginCount, 1)
		return token, nil
	}
}

// 探针确认会话仍有效时，刷新不应提交密码。
func TestRefreshTokenUsesProbeWithoutLogin(t *testing.T) {
	var loginCount int32
	r, store := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, acc config.Account) (string, error) {
		if acc.Token != "session-token" {
			t.Fatalf("probe should see current token, got %q", acc.Token)
		}
		return "session-token", nil
	}
	a := newProbeTestAuth("session-token")

	if !r.RefreshToken(context.Background(), a) {
		t.Fatal("expected refresh to succeed via session probe")
	}
	if got := atomic.LoadInt32(&loginCount); got != 0 {
		t.Fatalf("expected no password login, got %d", got)
	}
	if a.DeepSeekToken != "session-token" {
		t.Fatalf("expected existing token to be kept, got %q", a.DeepSeekToken)
	}
	if token := findTestToken(t, store); token != "session-token" {
		t.Fatalf("expected persisted token unchanged, got %q", token)
	}
}

// 服务端主动轮换时，探针拿回的新 token 要落盘并生效。
func TestRefreshTokenAdoptsRotatedTokenFromProbe(t *testing.T) {
	var loginCount int32
	r, store := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "rotated-token", nil
	}
	a := newProbeTestAuth("session-token")

	if !r.RefreshToken(context.Background(), a) {
		t.Fatal("expected refresh to succeed via session probe")
	}
	if got := atomic.LoadInt32(&loginCount); got != 0 {
		t.Fatalf("expected no password login, got %d", got)
	}
	if a.DeepSeekToken != "rotated-token" {
		t.Fatalf("expected rotated token on auth, got %q", a.DeepSeekToken)
	}
	if token := findTestToken(t, store); token != "rotated-token" {
		t.Fatalf("expected rotated token persisted, got %q", token)
	}
}

// 探针明确判定会话失效时才回退到密码登录。
func TestRefreshTokenFallsBackToLoginWhenSessionRejected(t *testing.T) {
	var loginCount int32
	r, store := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "", fmt.Errorf("check session: %w", ErrSessionExpired)
	}
	a := newProbeTestAuth("session-token")

	if !r.RefreshToken(context.Background(), a) {
		t.Fatal("expected fallback login to succeed")
	}
	if got := atomic.LoadInt32(&loginCount); got != 1 {
		t.Fatalf("expected exactly one login, got %d", got)
	}
	if a.DeepSeekToken != "login-token" {
		t.Fatalf("expected login token, got %q", a.DeepSeekToken)
	}
	if token := findTestToken(t, store); token != "login-token" {
		t.Fatalf("expected login token persisted, got %q", token)
	}
}

// 探针无法判定时仍保留密码登录兜底，避免把可恢复的请求判死。
func TestRefreshTokenFallsBackToLoginWhenProbeInconclusive(t *testing.T) {
	var loginCount int32
	r, _ := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "", errors.New("dial tcp: connection refused")
	}
	a := newProbeTestAuth("session-token")

	if !r.RefreshToken(context.Background(), a) {
		t.Fatal("expected fallback login to succeed")
	}
	if got := atomic.LoadInt32(&loginCount); got != 1 {
		t.Fatalf("expected exactly one login, got %d", got)
	}
	if a.DeepSeekToken != "login-token" {
		t.Fatalf("expected login token, got %q", a.DeepSeekToken)
	}
}

// 定时体检：探针成功即完成，不应触发登录。
func TestScheduledRefreshUsesProbeWithoutLogin(t *testing.T) {
	var loginCount int32
	r, _ := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "session-token", nil
	}
	markRefreshDue(t, r)
	a := newProbeTestAuth("session-token")

	if err := r.ensureManagedToken(context.Background(), a); err != nil {
		t.Fatalf("ensure managed token failed: %v", err)
	}
	if got := atomic.LoadInt32(&loginCount); got != 0 {
		t.Fatalf("expected no password login, got %d", got)
	}
	if a.DeepSeekToken != "session-token" {
		t.Fatalf("expected token kept, got %q", a.DeepSeekToken)
	}
}

// 定时体检：探针结果不可判定时保留 token，不能因为一次网络抖动就提交密码。
func TestScheduledRefreshKeepsTokenWhenProbeInconclusive(t *testing.T) {
	var loginCount int32
	r, _ := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "", errors.New("connection reset by peer")
	}
	markRefreshDue(t, r)
	a := newProbeTestAuth("session-token")

	if err := r.ensureManagedToken(context.Background(), a); err != nil {
		t.Fatalf("ensure managed token failed: %v", err)
	}
	if got := atomic.LoadInt32(&loginCount); got != 0 {
		t.Fatalf("expected no password login on inconclusive probe, got %d", got)
	}
	if a.DeepSeekToken != "session-token" {
		t.Fatalf("expected token kept, got %q", a.DeepSeekToken)
	}
}

// 定时体检：探针明确判定失效时回退到密码登录。
func TestScheduledRefreshLogsInWhenSessionRejected(t *testing.T) {
	var loginCount int32
	r, _ := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	r.RefreshSession = func(_ context.Context, _ config.Account) (string, error) {
		return "", ErrSessionExpired
	}
	markRefreshDue(t, r)
	a := newProbeTestAuth("session-token")

	if err := r.ensureManagedToken(context.Background(), a); err != nil {
		t.Fatalf("ensure managed token failed: %v", err)
	}
	if got := atomic.LoadInt32(&loginCount); got != 1 {
		t.Fatalf("expected exactly one login, got %d", got)
	}
	if a.DeepSeekToken != "login-token" {
		t.Fatalf("expected login token, got %q", a.DeepSeekToken)
	}
}

// 未注入探针时保持历史行为：定时到期直接登录。
func TestScheduledRefreshWithoutProbeFallsBackToLogin(t *testing.T) {
	var loginCount int32
	r, _ := newProbeTestResolver(t, countingLogin(&loginCount, "login-token"))
	markRefreshDue(t, r)
	a := newProbeTestAuth("session-token")

	if err := r.ensureManagedToken(context.Background(), a); err != nil {
		t.Fatalf("ensure managed token failed: %v", err)
	}
	if got := atomic.LoadInt32(&loginCount); got != 1 {
		t.Fatalf("expected exactly one login, got %d", got)
	}
}

func markRefreshDue(t *testing.T, r *Resolver) {
	t.Helper()
	if err := r.Store.Update(func(c *config.Config) error {
		c.Runtime.TokenRefreshIntervalHours = 1
		return nil
	}); err != nil {
		t.Fatalf("update runtime failed: %v", err)
	}
	r.mu.Lock()
	r.tokenRefreshedAt["acc@example.com"] = time.Now().Add(-2 * time.Hour)
	r.mu.Unlock()
}

func findTestToken(t *testing.T, store *config.Store) string {
	t.Helper()
	acc, ok := store.FindAccount("acc@example.com")
	if !ok {
		t.Fatal("account not found in store")
	}
	return acc.Token
}