111 lines
3.0 KiB
Go
111 lines
3.0 KiB
Go
package server
|
|
|
|
import (
|
|
"bufio"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"hunjing/shenshi/internal/route"
|
|
)
|
|
|
|
func newTestServer(token string) *httptest.Server {
|
|
srv := New(Config{ServiceToken: token}, route.New())
|
|
return httptest.NewServer(srv)
|
|
}
|
|
|
|
func chatBody() string {
|
|
return `{"ir_version":1,"slot":"chat","messages":[` +
|
|
`{"role":"system","content":"契约占位"},{"role":"user","content":"你好"}],` +
|
|
`"run_id":"run_test","stream":true}`
|
|
}
|
|
|
|
func post(t *testing.T, ts *httptest.Server, token, body string) (*http.Response, string) {
|
|
t.Helper()
|
|
req, err := http.NewRequest(http.MethodPost, ts.URL+"/v1/chat/completions", strings.NewReader(body))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if token != "" {
|
|
req.Header.Set("Authorization", "Bearer "+token)
|
|
}
|
|
resp, err := ts.Client().Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Cleanup(func() { _ = resp.Body.Close() })
|
|
var sb strings.Builder
|
|
scanner := bufio.NewScanner(resp.Body)
|
|
for scanner.Scan() {
|
|
sb.WriteString(scanner.Text())
|
|
sb.WriteString("\n")
|
|
}
|
|
return resp, sb.String()
|
|
}
|
|
|
|
func TestChatCompletionsStreamsEchoAndDone(t *testing.T) {
|
|
ts := newTestServer("")
|
|
resp, body := post(t, ts, "", chatBody())
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
|
}
|
|
if got := resp.Header.Get("X-Accel-Buffering"); got != "no" {
|
|
t.Fatalf("X-Accel-Buffering = %q", got)
|
|
}
|
|
if !strings.Contains(body, `"delta":"【神识回声桩】"`) && !strings.Contains(body, "神识回声桩") {
|
|
t.Fatalf("缺少回声增量: %s", body)
|
|
}
|
|
if !strings.Contains(body, `"finish_reason":"stop"`) {
|
|
t.Fatalf("缺少终帧用量: %s", body)
|
|
}
|
|
if !strings.Contains(body, "data: [DONE]") {
|
|
t.Fatalf("缺少 [DONE]: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestServiceTokenRejectsInvalidCredential(t *testing.T) {
|
|
ts := newTestServer("secret-token")
|
|
resp, body := post(t, ts, "wrong", chatBody())
|
|
if resp.StatusCode != http.StatusUnauthorized {
|
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
|
}
|
|
if !strings.Contains(body, "HJ-2003") {
|
|
t.Fatalf("缺少错误码: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestInvalidIRRejected(t *testing.T) {
|
|
ts := newTestServer("")
|
|
resp, body := post(t, ts, "", `{"ir_version":99,"messages":[],"run_id":"","stream":true}`)
|
|
if resp.StatusCode != http.StatusBadRequest {
|
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
|
}
|
|
if !strings.Contains(body, "HJ-2002") {
|
|
t.Fatalf("缺少错误码: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestUnknownSlotNoEndpoint(t *testing.T) {
|
|
ts := newTestServer("")
|
|
resp, body := post(t, ts, "", `{"ir_version":1,"slot":"vision","messages":[{"role":"user","content":"x"}],"run_id":"run_t","stream":true}`)
|
|
if resp.StatusCode != http.StatusServiceUnavailable {
|
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
|
}
|
|
if !strings.Contains(body, "HJ-2001") {
|
|
t.Fatalf("缺少错误码: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestHealthz(t *testing.T) {
|
|
ts := newTestServer("")
|
|
resp, err := ts.Client().Get(ts.URL + "/healthz")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
if resp.StatusCode != http.StatusOK {
|
|
t.Fatalf("状态码 = %d", resp.StatusCode)
|
|
}
|
|
}
|