package server import ( "bytes" "context" "encoding/json" "fmt" "io" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "git.misaka.ren/M1saka/zhanlu_proxy/internal/auth" "git.misaka.ren/M1saka/zhanlu_proxy/internal/config" ) const testSM2Key = "8d6ee90b3c4d299ae5abd655dbc3547c110ae8aeff1de18b0df241f215f90748" // setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to it. func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, string) { t.Helper() upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/query/acepilot-h5/manager/code/getAuthCode": writeJSON(w, http.StatusOK, map[string]any{"state": "OK"}) case "/api/query/acepilot-h5/manager/code/checkCode": writeJSON(w, http.StatusOK, map[string]any{"state": "OK", "body": map[string]any{ "result": true, "ak": "BASE64AK", "sk": "BASE64SK", "license": "BASE64TOKEN", }}) case "/api/acepilot/zhanlu/v1/login": if r.Header.Get("plugin_type") != "zhanlu_ide" { writeJSON(w, http.StatusBadRequest, map[string]any{"state": "ERROR", "errorMessage": "bad plugin_type"}) return } writeJSON(w, http.StatusOK, map[string]any{"state": "OK", "body": map[string]any{ "email": "dev@example.com", "organization": "cmcc", "team": "ai", }}) case "/user/api/v2/external/key/get-or-create": writeJSON(w, http.StatusOK, map[string]any{"apiKey": "sk-test-456"}) case "/chat/completions": if r.Header.Get("Authorization") != "Bearer sk-test-456" { writeJSON(w, http.StatusUnauthorized, map[string]any{"error": map[string]any{"message": "bad auth"}}) return } w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":1}}\n\n") _, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") _, _ = fmt.Fprint(w, "data: [DONE]\n\n") case "/gateway/v1/model/info": writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{ {"model_name": "GLM-4.7"}, {"id": "MiniMaxAI/MiniMax-M2.5"}, }}) default: http.NotFound(w, r) } })) credsFile := filepath.Join(t.TempDir(), "credentials.json") cfg := config.Config{ ListenAddr: ":0", MobileLoginBaseURL: upstream.URL, MobileModelBaseURL: upstream.URL, UpstreamPath: "/chat/completions", CredentialsPath: credsFile, TokenDecryptKey: "3jw7woww2rvhla6k", PublicKeyPEM: defaultTestPublicKey, PhonePublicKeyPEM: defaultTestPublicKey, SM2PrivateKey: testSM2Key, PluginVersion: "1.4.2", } h := New(cfg) proxy := httptest.NewServer(h) return upstream, proxy, credsFile } // TestPhoneLoginAndChat exercises the full v1.4.2 flow: SMS login, profile // fetch, SM2 API-key provisioning, then OpenAI-compatible chat and models. func TestPhoneLoginAndChat(t *testing.T) { upstream, proxy, credsFile := setupTestServer(t) defer upstream.Close() defer proxy.Close() // 1. request phone code resp, err := http.Post(proxy.URL+"/api/auth/code", "application/json", strings.NewReader(`{"telephone":"13800000000"}`)) if err != nil { t.Fatal(err) } var codeResp map[string]any _ = json.NewDecoder(resp.Body).Decode(&codeResp) resp.Body.Close() secret, _ := codeResp["secret"].(string) // 2. login with phone code (server RSA-encrypts the telephone itself) loginBody, _ := json.Marshal(map[string]string{"telephone": "13800000000", "code": "123456", "secret": secret}) resp, err = http.Post(proxy.URL+"/api/auth/login", "application/json", bytes.NewReader(loginBody)) if err != nil { t.Fatal(err) } var loginResp map[string]any _ = json.NewDecoder(resp.Body).Decode(&loginResp) resp.Body.Close() if !okValue(loginResp) { t.Fatalf("login failed: %v", loginResp) } // 3. credentials file should contain the provisioned api key creds, err := auth.LoadCredentials(credsFile) if err != nil { t.Fatal(err) } if creds.APIKey != "sk-test-456" { t.Fatalf("apiKey = %q", creds.APIKey) } if creds.Email != "dev@example.com" { t.Fatalf("email = %q", creds.Email) } // 4. non-streaming chat completion chatBody := `{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}` resp, err = http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(chatBody)) if err != nil { t.Fatal(err) } defer resp.Body.Close() var chatResp map[string]any if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil { t.Fatal(err) } choices, _ := chatResp["choices"].([]any) if len(choices) != 1 { t.Fatalf("chat choices = %v", chatResp) } msg, _ := choices[0].(map[string]any)["message"].(map[string]any) if msg["content"] != "hello" { t.Fatalf("chat content = %v", msg) } // 5. models endpoint should prefer gateway model info resp, err = http.Get(proxy.URL + "/v1/models") if err != nil { t.Fatal(err) } defer resp.Body.Close() var modelsResp map[string]any _ = json.NewDecoder(resp.Body).Decode(&modelsResp) items, _ := modelsResp["data"].([]any) if len(items) != 2 { t.Fatalf("models = %v", modelsResp) } } func TestStreamingChat(t *testing.T) { upstream, proxy, credsFile := setupTestServer(t) defer upstream.Close() defer proxy.Close() // Seed credentials directly with the api key creds := auth.Credentials{ AccessKey: "AK", SecretKey: "SK", Token: "TOKEN", APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "dev@example.com", } if err := auth.SaveCredentials(credsFile, creds); err != nil { t.Fatal(err) } chatBody := `{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":true}` resp, err := http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(chatBody)) if err != nil { t.Fatal(err) } defer resp.Body.Close() if resp.StatusCode != 200 { b, _ := io.ReadAll(resp.Body) t.Fatalf("status %d: %s", resp.StatusCode, string(b)) } raw, _ := io.ReadAll(resp.Body) body := string(raw) if !strings.Contains(body, "data: ") || !strings.Contains(body, "hello") || !strings.Contains(body, "[DONE]") { t.Fatalf("stream body: %s", body) } } func okValue(m map[string]any) bool { ok, _ := m["ok"].(bool) return ok } func TestMain(m *testing.M) { os.Exit(m.Run()) } var _ = context.Background const defaultTestPublicKey = `-----BEGIN PUBLIC KEY----- MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAhxudxTewPgljUHEZHkusP7m3I+zA4/RGvuUMt6TtII/m4zwUOm/Y31zHBTmkCCt8k5vj9y+AmO0TsGmHooNQuMebakdmEWdcA5h7YAHHFbF2w5LcxIXjib08vgVpA+m3R5xPbLK+vfHe2aAX36b5nHReDNncY5vAl3U4CgIEBGPqyG67vJytRWqP+sfEdw5+m192Rf4SCGyiBzRmjiVlH3zeEBjdbOrkAnzKOVz6AHBl2q7LPLJKIzxjoAyhEp5qnDjHUFo5VZUgFwUOt83A/jbGMyzmjRoxBuvKcs9tBuorZyUwIsZN6E+rtQk2YqMPj4RkDsZ7LRmj6on8sN2rHQIDAQAB -----END PUBLIC KEY-----`