internal/upstream/device_token_test.go

// device_token_test.go X-Device-Token 注入 + 文件兜底读取单测。
package upstream

import (
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/linguo2625469/workbuddy2api-panel/internal/auth"
)

// TestDeviceTokenInjected_WhenSet auth.Auth.DeviceToken 非空时 chat/billing 请求均注入。
func TestDeviceTokenInjected_WhenSet(t *testing.T) {
	a := &auth.Auth{AccessToken: "at", UID: "u1", DeviceToken: "tok-from-auth"}
	// 用 server 端验证而非 RoundTripper 捕获：更贴近真实注入路径。
	for _, tc := range []struct {
		name     string
		apply    func(c *Client, req *http.Request)
		wantPath string
	}{
		{"chat", func(c *Client, req *http.Request) { c.ChatHeaders(req, a, "", ChatMeta{}) }, "/v2/chat/completions"},
		{"billing", func(c *Client, req *http.Request) { c.BillingHeaders(req, a) }, "/v2/report"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var got string
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				got = r.Header.Get("X-Device-Token")
				w.WriteHeader(200)
				_, _ = w.Write([]byte(`{"code":0}`))
			}))
			defer srv.Close()
			c := &Client{
				HTTP:          srv.Client(),
				ChatHTTP:      srv.Client(),
				ChatBaseCN:    srv.URL,
				BillingBaseCN: srv.URL,
			}
			req, _ := http.NewRequest(http.MethodPost, srv.URL+tc.wantPath, nil)
			tc.apply(c, req)
			resp, err := c.HTTP.Do(req)
			if err != nil {
				t.Fatalf("do: %v", err)
			}
			resp.Body.Close()
			if got != "tok-from-auth" {
				t.Errorf("X-Device-Token = %q want %q", got, "tok-from-auth")
			}
		})
	}
}

// TestDeviceTokenNotInjected_WhenEmpty 所有来源皆空时不注入该头。
func TestDeviceTokenNotInjected_WhenEmpty(t *testing.T) {
	a := &auth.Auth{AccessToken: "at", UID: "u1"} // DeviceToken 空
	var got string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got = r.Header.Get("X-Device-Token")
		w.WriteHeader(200)
		_, _ = w.Write([]byte(`{"code":0}`))
	}))
	defer srv.Close()
	c := &Client{
		HTTP:          srv.Client(),
		ChatHTTP:      srv.Client(),
		ChatBaseCN:    srv.URL,
		BillingBaseCN: srv.URL,
		// DeviceToken / DeviceTokenFile 皆空
	}
	req, _ := http.NewRequest(http.MethodPost, srv.URL+"/v2/chat/completions", nil)
	c.ChatHeaders(req, a, "", ChatMeta{})
	resp, err := c.HTTP.Do(req)
	if err != nil {
		t.Fatalf("do: %v", err)
	}
	resp.Body.Close()
	if got != "" {
		t.Errorf("X-Device-Token = %q want empty (not injected)", got)
	}
}

// TestDeviceTokenFromConfigOrFile_Overrides 优先级：auth > config > 文件。
// auth 有值时覆盖 config；auth 空时 config 兜底；auth 与 config 皆空时读文件。
func TestDeviceTokenFromConfigOrFile_Overrides(t *testing.T) {
	dir := t.TempDir()
	fp := filepath.Join(dir, "device_token")
	if err := os.WriteFile(fp, []byte("tok-from-file\n"), 0o600); err != nil {
		t.Fatalf("write file: %v", err)
	}

	cases := []struct {
		name string
		auth string
		cfg  string
		want string
	}{
		{"auth_over_config", "tok-auth", "tok-config", "tok-auth"},
		{"config_when_auth_empty", "", "tok-config", "tok-config"},
		{"file_when_auth_and_config_empty", "", "", "tok-from-file"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			// 每个用例用独立缓存：device token 文件缓存 5 分钟，case 间会串扰。
			save := resetDeviceTokenFileCache(fp)
			defer save()
			a := &auth.Auth{AccessToken: "at", UID: "u1", DeviceToken: tc.auth}
			var got string
			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				got = r.Header.Get("X-Device-Token")
				w.WriteHeader(200)
				_, _ = w.Write([]byte(`{"code":0}`))
			}))
			defer srv.Close()
			c := &Client{
				HTTP:            srv.Client(),
				ChatHTTP:        srv.Client(),
				ChatBaseCN:      srv.URL,
				BillingBaseCN:   srv.URL,
				DeviceToken:     tc.cfg,
				DeviceTokenFile: fp,
			}
			req, _ := http.NewRequest(http.MethodPost, srv.URL+"/v2/chat/completions", nil)
			c.ChatHeaders(req, a, "", ChatMeta{})
			resp, err := c.HTTP.Do(req)
			if err != nil {
				t.Fatalf("do: %v", err)
			}
			resp.Body.Close()
			if got != tc.want {
				t.Errorf("X-Device-Token = %q want %q", got, tc.want)
			}
		})
	}
}

// TestDeviceTokenFileTooLarge 文件超过 1KB 时忽略不注入。
func TestDeviceTokenFileTooLarge(t *testing.T) {
	dir := t.TempDir()
	fp := filepath.Join(dir, "device_token")
	if err := os.WriteFile(fp, []byte(strings.Repeat("x", 2048)), 0o600); err != nil {
		t.Fatalf("write file: %v", err)
	}
	save := resetDeviceTokenFileCache(fp)
	defer save()
	a := &auth.Auth{AccessToken: "at", UID: "u1"}
	c := &Client{DeviceTokenFile: fp}
	if tok := c.resolveDeviceToken(a); tok != "" {
		t.Errorf("resolveDeviceToken() = %q want empty (file too large)", tok)
	}
}

// resetDeviceTokenFileCache 替换全局 device token 文件缓存并返回恢复函数。
// 文件缓存 5 分钟 TTL，测试间需清空避免串扰。
func resetDeviceTokenFileCache(path string) (restore func()) {
	dtFileCache.mu.Lock()
	origPath := dtFileCache.path
	origTok := dtFileCache.token
	origRead := dtFileCache.readAt
	origErr := dtFileCache.lastErr
	dtFileCache.path = path
	dtFileCache.token = ""
	dtFileCache.readAt = time.Time{}
	dtFileCache.lastErr = nil
	dtFileCache.mu.Unlock()
	var once sync.Once
	return func() {
		once.Do(func() {
			dtFileCache.mu.Lock()
			dtFileCache.path = origPath
			dtFileCache.token = origTok
			dtFileCache.readAt = origRead
			dtFileCache.lastErr = origErr
			dtFileCache.mu.Unlock()
		})
	}
}