zhanlu_proxy/internal/zhanlu/client_test.go

package zhanlu

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/sign"
)

const testSM2Key = "8d6ee90b3c4d299ae5abd655dbc3547c110ae8aeff1de18b0df241f215f90748"

func TestClientFlow(t *testing.T) {
	var gotLoginHeaders, gotProvisionHeaders http.Header
	var gotChatAuth, gotModelsAuth string

	mux := http.NewServeMux()
	mux.HandleFunc("/api/acepilot/zhanlu/v1/login", func(w http.ResponseWriter, r *http.Request) {
		gotLoginHeaders = r.Header
		// echo a profile whose fields are plaintext (DecryptCredentialOrRaw fallback)
		writeJSON(t, w, map[string]any{"state": "OK", "body": map[string]any{
			"email": "[email protected]", "organization": "cmcc", "team": "ai", "name": "Dev", "telephone": "13800000000",
		}})
	})
	mux.HandleFunc("/user/api/v2/external/key/get-or-create", func(w http.ResponseWriter, r *http.Request) {
		gotProvisionHeaders = r.Header
		if gotProvisionHeaders.Get("X-Auth-Signature") == "" || gotProvisionHeaders.Get("X-Auth-Timestamp") == "" || gotProvisionHeaders.Get("X-Auth-Nonce") == "" {
			t.Errorf("provision request missing X-Auth-* headers: %v", gotProvisionHeaders)
		}
		if gotProvisionHeaders.Get("X-Auth-Nonce") == "" || len(gotProvisionHeaders.Get("X-Auth-Nonce")) != 32 {
			t.Errorf("X-Auth-Nonce should be 32 chars")
		}
		writeJSON(t, w, map[string]any{"apiKey": "sk-zhanlu-test-123"})
	})
	mux.HandleFunc("/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		gotChatAuth = r.Header.Get("Authorization")
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-1\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")
		_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
	})
	mux.HandleFunc("/gateway/v1/model/info", func(w http.ResponseWriter, r *http.Request) {
		gotModelsAuth = r.Header.Get("Authorization")
		writeJSON(t, w, map[string]any{"data": []map[string]any{
			{"model_name": "GLM-4.7"},
			{"id": "MiniMaxAI/MiniMax-M2.5"},
			{"model_info": map[string]any{"id": "qwen-max"}},
		}})
	})

	ts := httptest.NewServer(mux)
	defer ts.Close()

	pub, err := auth.ParsePublicKey(`-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAhxudxTewPgljUHEZHkusP7m3I+zA4/RGvuUMt6TtII/m4zwUOm/Y31zHBTmkCCt8k5vj9y+AmO0TsGmHooNQuMebakdmEWdcA5h7YAHHFbF2w5LcxIXjib08vgVpA+m3R5xPbLK+vfHe2aAX36b5nHReDNncY5vAl3U4CgIEBGPqyG67vJytRWqP+sfEdw5+m192Rf4SCGyiBzRmjiVlH3zeEBjdbOrkAnzKOVz6AHBl2q7LPLJKIzxjoAyhEp5qnDjHUFo5VZUgFwUOt83A/jbGMyzmjRoxBuvKcs9tBuorZyUwIsZN6E+rtQk2YqMPj4RkDsZ7LRmj6on8sN2rHQIDAQAB
-----END PUBLIC KEY-----`)
	if err != nil {
		t.Fatal(err)
	}
	c := NewClient(ts.URL, ts.URL, "/chat/completions", "1.4.2", testSM2Key, sign.Signer{PublicKey: pub}, 0)

	creds := auth.Credentials{AccessKey: "AK", SecretKey: "SK", Token: "TOKEN"}
	profile, err := c.LoginProfile(context.Background(), creds)
	if err != nil {
		t.Fatalf("LoginProfile: %v", err)
	}
	if profile.Email != "[email protected]" {
		t.Fatalf("profile email = %q", profile.Email)
	}
	if gotLoginHeaders.Get("plugin_type") != "zhanlu_ide" {
		t.Errorf("plugin_type header = %q", gotLoginHeaders.Get("plugin_type"))
	}
	if gotLoginHeaders.Get("plugin_version") != "1.4.2" {
		t.Errorf("plugin_version header = %q", gotLoginHeaders.Get("plugin_version"))
	}

	apiKey, err := c.ProvisionAPIKey(context.Background(), profile.Email, profile.Organization, profile.Team)
	if err != nil {
		t.Fatalf("ProvisionAPIKey: %v", err)
	}
	if apiKey != "sk-zhanlu-test-123" {
		t.Fatalf("apiKey = %q", apiKey)
	}

	resp, err := c.ChatCompletions(context.Background(), apiKey, []byte(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}]}`))
	if err != nil {
		t.Fatalf("ChatCompletions: %v", err)
	}
	defer resp.Body.Close()
	if gotChatAuth != "Bearer sk-zhanlu-test-123" {
		t.Errorf("chat Authorization = %q", gotChatAuth)
	}

	models, err := c.Models(context.Background(), apiKey)
	if err != nil {
		t.Fatalf("Models: %v", err)
	}
	if len(models) != 3 || models[0] != "GLM-4.7" || models[1] != "MiniMaxAI/MiniMax-M2.5" || models[2] != "qwen-max" {
		t.Fatalf("models = %v", models)
	}
	if gotModelsAuth != "Bearer sk-zhanlu-test-123" {
		t.Errorf("models Authorization = %q", gotModelsAuth)
	}
}

func TestProvisionAPIKeyDefaultsPlaceholders(t *testing.T) {
	var body string
	mux := http.NewServeMux()
	mux.HandleFunc("/user/api/v2/external/key/get-or-create", func(w http.ResponseWriter, r *http.Request) {
		buf := make([]byte, 512)
		n, _ := r.Body.Read(buf)
		body = strings.TrimSpace(string(buf[:n]))
		writeJSON(t, w, map[string]any{"data": map[string]any{"key": "k2"}})
	})
	ts := httptest.NewServer(mux)
	defer ts.Close()

	c := NewClient(ts.URL, ts.URL, "/chat/completions", "1.4.2", testSM2Key, sign.Signer{}, 0)
	apiKey, err := c.ProvisionAPIKey(context.Background(), "[email protected]", "", "")
	if err != nil {
		t.Fatalf("ProvisionAPIKey: %v", err)
	}
	if apiKey != "k2" {
		t.Fatalf("apiKey = %q", apiKey)
	}
	var parsed map[string]string
	if err := json.Unmarshal([]byte(body), &parsed); err != nil {
		t.Fatal(err)
	}
	if parsed["organization"] != "未配置" || parsed["team"] != "未配置" {
		t.Fatalf("placeholders not applied: %v", parsed)
	}
}

func writeJSON(t *testing.T, w http.ResponseWriter, v any) {
	t.Helper()
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(v)
}