zhanlu_proxy/internal/server/server_test.go

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": "[email protected]", "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,
		Models:             []string{"GLM-4.7", "MiniMaxAI/MiniMax-M2.5"},
		DefaultModel:       "GLM-4.7",
		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 != "[email protected]" {
		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:        "[email protected]",
	}
	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-----`