397 lines
16 KiB
Go
397 lines
16 KiB
Go
// k3r053n3_test: Test suite for Kimi K3 Gateway
|
|
// Created by Luxferre in 2026, released into the public domain
|
|
|
|
package main
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestModelsEndpoint(t *testing.T) {
|
|
svc := NewKimiService("https://test.hf.space", "kimi-k3", "direct:together", "max", "", 8192, 0.7)
|
|
req, w := httptest.NewRequest("GET", "/v1/models", nil), httptest.NewRecorder()
|
|
NewMux(svc).ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected status 200, got %d", w.Code)
|
|
}
|
|
var res ModelsResponse
|
|
if err := json.NewDecoder(w.Body).Decode(&res); err != nil || len(res.Data) == 0 {
|
|
t.Fatalf("invalid models response: %v", err)
|
|
}
|
|
for _, m := range res.Data {
|
|
if m.ID == "kimi-k3" {
|
|
return
|
|
}
|
|
}
|
|
t.Fatalf("expected kimi-k3 in models list")
|
|
}
|
|
|
|
func TestResolveReasoningEffort(t *testing.T) {
|
|
tests := []struct{ in, def, exp string }{
|
|
{"", "max", "max"}, {"max", "default", "max"}, {"HIGH", "low", "high"},
|
|
{"low", "max", "low"}, {"medium", "max", "default"}, {"default", "max", "default"}, {"unknown", "max", "max"},
|
|
}
|
|
for _, tt := range tests {
|
|
if res := ResolveReasoningEffort(ChatCompletionRequest{ReasoningEffort: tt.in}, tt.def); res != tt.exp {
|
|
t.Errorf("ResolveReasoningEffort(%q, %q) = %q, want %q", tt.in, tt.def, res, tt.exp)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestResolveBackend(t *testing.T) {
|
|
svc := NewKimiService("https://test.hf.space", "kimi-k3", "direct:together", "max", "", 8192, 0.7)
|
|
tests := []struct{ model, exp string }{
|
|
{"kimi-k3", "direct:together"}, {"kimi-k3:fireworks", "direct:fireworks"},
|
|
{"kimi-k3:together", "direct:together"}, {"kimi-k3:hf-together", "hf:together"},
|
|
{"kimi-k3:hf-fireworks", "hf:fireworks-ai"}, {"kimi-k3:hf-featherless", "hf:featherless-ai"},
|
|
{"kimi-k3:hf-baseten", "hf:baseten"}, {"custom-model", "direct:together"},
|
|
}
|
|
for _, tt := range tests {
|
|
if b := svc.ResolveBackend(tt.model); b != tt.exp {
|
|
t.Errorf("ResolveBackend(%q) = %q, want %q", tt.model, b, tt.exp)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestParseSOCKS5URL(t *testing.T) {
|
|
tests := []struct {
|
|
in string
|
|
exp *SOCKS5Config
|
|
}{
|
|
{"", nil},
|
|
{"127.0.0.1:1080", &SOCKS5Config{Address: "127.0.0.1:1080"}},
|
|
{"socks5://127.0.0.1:9050", &SOCKS5Config{Address: "127.0.0.1:9050"}},
|
|
{"socks5h://user:pass@proxy.example.com:1080", &SOCKS5Config{Address: "proxy.example.com:1080", Username: "user", Password: "pass"}},
|
|
{"socks5://myuser@10.0.0.1:1080", &SOCKS5Config{Address: "10.0.0.1:1080", Username: "myuser"}},
|
|
{"127.0.0.1", &SOCKS5Config{Address: "127.0.0.1:1080"}},
|
|
}
|
|
for _, tt := range tests {
|
|
cfg, err := ParseSOCKS5URL(tt.in)
|
|
if err != nil {
|
|
t.Errorf("ParseSOCKS5URL(%q) err: %v", tt.in, err)
|
|
} else if tt.exp == nil && cfg != nil {
|
|
t.Errorf("ParseSOCKS5URL(%q) expected nil, got %+v", tt.in, cfg)
|
|
} else if tt.exp != nil && (cfg.Address != tt.exp.Address || cfg.Username != tt.exp.Username || cfg.Password != tt.exp.Password) {
|
|
t.Errorf("ParseSOCKS5URL(%q) = %+v, want %+v", tt.in, cfg, tt.exp)
|
|
}
|
|
}
|
|
}
|
|
|
|
func startMockSOCKS5Server(t *testing.T, expectedUser, expectedPass string) (string, func()) {
|
|
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("failed to start mock socks5 listener: %v", err)
|
|
}
|
|
stopCh := make(chan struct{})
|
|
|
|
go func() {
|
|
for {
|
|
conn, err := ln.Accept()
|
|
if err != nil {
|
|
return
|
|
}
|
|
go func(c net.Conn) {
|
|
defer c.Close()
|
|
hdr := make([]byte, 2)
|
|
if _, err := io.ReadFull(c, hdr); err != nil || hdr[0] != 0x05 {
|
|
return
|
|
}
|
|
methods := make([]byte, int(hdr[1]))
|
|
if _, err := io.ReadFull(c, methods); err != nil {
|
|
return
|
|
}
|
|
|
|
if expectedUser != "" {
|
|
c.Write([]byte{0x05, 0x02})
|
|
authHdr := make([]byte, 2)
|
|
if _, err := io.ReadFull(c, authHdr); err != nil {
|
|
return
|
|
}
|
|
userBuf := make([]byte, int(authHdr[1]))
|
|
io.ReadFull(c, userBuf)
|
|
pLenBuf := make([]byte, 1)
|
|
io.ReadFull(c, pLenBuf)
|
|
passBuf := make([]byte, int(pLenBuf[0]))
|
|
io.ReadFull(c, passBuf)
|
|
|
|
if string(userBuf) != expectedUser || string(passBuf) != expectedPass {
|
|
c.Write([]byte{0x01, 0x01})
|
|
return
|
|
}
|
|
c.Write([]byte{0x01, 0x00})
|
|
} else {
|
|
c.Write([]byte{0x05, 0x00})
|
|
}
|
|
|
|
reqHdr := make([]byte, 4)
|
|
if _, err := io.ReadFull(c, reqHdr); err != nil || reqHdr[0] != 0x05 || reqHdr[1] != 0x01 {
|
|
return
|
|
}
|
|
var targetAddr string
|
|
switch reqHdr[3] {
|
|
case 0x01:
|
|
ipBuf := make([]byte, 4)
|
|
io.ReadFull(c, ipBuf)
|
|
targetAddr = net.IP(ipBuf).String()
|
|
case 0x03:
|
|
lBuf := make([]byte, 1)
|
|
io.ReadFull(c, lBuf)
|
|
dBuf := make([]byte, int(lBuf[0]))
|
|
io.ReadFull(c, dBuf)
|
|
targetAddr = string(dBuf)
|
|
}
|
|
portBuf := make([]byte, 2)
|
|
io.ReadFull(c, portBuf)
|
|
port := (int(portBuf[0]) << 8) | int(portBuf[1])
|
|
|
|
remote, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", targetAddr, port), 5*time.Second)
|
|
if err != nil {
|
|
c.Write([]byte{0x05, 0x05, 0x00, 0x01, 0, 0, 0, 0, 0, 0})
|
|
return
|
|
}
|
|
defer remote.Close()
|
|
c.Write([]byte{0x05, 0x00, 0x00, 0x01, 127, 0, 0, 1, 0x1f, 0x90})
|
|
go io.Copy(remote, c)
|
|
io.Copy(c, remote)
|
|
}(conn)
|
|
}
|
|
}()
|
|
return ln.Addr().String(), func() { close(stopCh); ln.Close() }
|
|
}
|
|
|
|
func TestSOCKS5EndToEnd(t *testing.T) {
|
|
mockTarget := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/test" {
|
|
w.Write([]byte("socks5-ok"))
|
|
return
|
|
}
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer mockTarget.Close()
|
|
|
|
targetHostPort := strings.TrimPrefix(mockTarget.URL, "http://")
|
|
proxyAddr, cleanupProxy := startMockSOCKS5Server(t, "alice", "secret123")
|
|
defer cleanupProxy()
|
|
|
|
proxyURL := fmt.Sprintf("socks5://alice:secret123@%s", proxyAddr)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
conn, err := DialSOCKS5(ctx, proxyURL, targetHostPort)
|
|
if err != nil {
|
|
t.Fatalf("DialSOCKS5 failed: %v", err)
|
|
}
|
|
defer conn.Close()
|
|
|
|
fmt.Fprintf(conn, "GET /test HTTP/1.1\r\nHost: %s\r\nConnection: close\r\n\r\n", targetHostPort)
|
|
respBody, err := io.ReadAll(conn)
|
|
if err != nil || !strings.Contains(string(respBody), "socks5-ok") {
|
|
t.Fatalf("socks5 direct dial failed: %v, got: %s", err, string(respBody))
|
|
}
|
|
|
|
svc := NewKimiService(mockTarget.URL, "kimi-k3", "direct:together", "max", proxyURL, 8192, 0.7)
|
|
req, _ := http.NewRequest("GET", mockTarget.URL+"/test", nil)
|
|
resp, err := svc.client.Do(req)
|
|
if err != nil {
|
|
t.Fatalf("svc.client via SOCKS5 failed: %v", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, _ := io.ReadAll(resp.Body)
|
|
if string(body) != "socks5-ok" {
|
|
t.Fatalf("expected 'socks5-ok', got: %s", string(body))
|
|
}
|
|
}
|
|
|
|
func TestParseGradioSnapshot(t *testing.T) {
|
|
// Test 1: Standard reasoning done and answer present with internal </div> tag in reasoning
|
|
sample1 := `[[
|
|
{"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}],"options":null},
|
|
{"role":"assistant","metadata":null,"content":[{"text":"<","type":"text"}],"options":null},
|
|
{"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":">I am thinking about </div> and </tag> tags.</","type":"text"}],"options":null},
|
|
{"role":"assistant","metadata":null,"content":[{"text":">\n\nHello there!","type":"text"}],"options":null}
|
|
], {"value":{"text":"","files":[]},"__type__":"update"}]`
|
|
|
|
content, reasoning, ok := parseGradioSnapshot(sample1)
|
|
if !ok || content != "Hello there!" || reasoning != "I am thinking about </div> and </tag> tags." {
|
|
t.Fatalf("parseGradioSnapshot sample1 failed: content=%q, reasoning=%q", content, reasoning)
|
|
}
|
|
|
|
// Test 2: Reasoning still pending
|
|
sample2 := `[[
|
|
{"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}],"options":null},
|
|
{"role":"assistant","metadata":null,"content":[{"text":"<","type":"text"}],"options":null},
|
|
{"role":"assistant","metadata":{"title":"Reasoning","status":"pending"},"content":[{"text":">Still reasoning","type":"text"}],"options":null}
|
|
], {"value":{"text":"","files":[]},"__type__":"update"}]`
|
|
|
|
content2, reasoning2, ok2 := parseGradioSnapshot(sample2)
|
|
if !ok2 || content2 != "" || reasoning2 != "Still reasoning" {
|
|
t.Fatalf("parseGradioSnapshot sample2 failed: content=%q, reasoning=%q", content2, reasoning2)
|
|
}
|
|
|
|
// Test 3: Interleaved messages
|
|
sample3 := `[[
|
|
{"role":"user","metadata":null,"content":[{"text":"Hi","type":"text"}]},
|
|
{"role":"assistant","metadata":null,"content":[{"text":"<"}]},
|
|
{"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":">Part 1"}]},
|
|
{"role":"assistant","metadata":null,"content":[{"text":"intro text: "}]},
|
|
{"role":"assistant","metadata":{"title":"Reasoning","status":"done"},"content":[{"text":". Part 2</"}]},
|
|
{"role":"assistant","metadata":null,"content":[{"text":">\n\n<tool_call>{\"name\":\"fn\"}</tool_call>"}]}
|
|
]]`
|
|
content3, reasoning3, ok3 := parseGradioSnapshot(sample3)
|
|
if !ok3 || !strings.Contains(content3, "<tool_call>") || !strings.Contains(reasoning3, "Part 1. Part 2") {
|
|
t.Fatalf("parseGradioSnapshot sample3 failed: content=%q, reasoning=%q", content3, reasoning3)
|
|
}
|
|
}
|
|
|
|
func TestDetectToolCalls(t *testing.T) {
|
|
// 1. Standard XML tag with conversational text
|
|
text1 := "Here is the weather:\n<tool_call>\n{\"name\": \"get_weather\", \"arguments\": {\"location\": \"Tokyo\"}}\n</tool_call>\nDone."
|
|
tcs1, cleaned1, ok1 := DetectToolCalls(text1)
|
|
if !ok1 || len(tcs1) != 1 || tcs1[0].Function.Name != "get_weather" || !strings.Contains(tcs1[0].Function.Arguments, "Tokyo") {
|
|
t.Fatalf("DetectToolCalls 1 failed: ok=%v, tcs=%+v", ok1, tcs1)
|
|
}
|
|
if strings.Contains(cleaned1, "<tool_call>") || !strings.Contains(cleaned1, "Here is the weather:") {
|
|
t.Fatalf("DetectToolCalls cleaned 1 invalid: %q", cleaned1)
|
|
}
|
|
|
|
// 2. tool-call variant
|
|
text2 := "<tool-call>\n{\"name\": \"calc\", \"arguments\": \"{\\\"x\\\": 1}\"}\n</tool-call>"
|
|
tcs2, _, ok2 := DetectToolCalls(text2)
|
|
if !ok2 || len(tcs2) != 1 || tcs2[0].Function.Name != "calc" {
|
|
t.Fatalf("DetectToolCalls 2 failed: ok=%v, tcs=%+v", ok2, tcs2)
|
|
}
|
|
|
|
// 3. markdown codeblock format
|
|
text3 := "```tool_call\n{\"name\": \"search\", \"query\": \"golang\"}\n```"
|
|
tcs3, _, ok3 := DetectToolCalls(text3)
|
|
if !ok3 || len(tcs3) != 1 || tcs3[0].Function.Name != "search" {
|
|
t.Fatalf("DetectToolCalls 3 failed: ok=%v, tcs=%+v", ok3, tcs3)
|
|
}
|
|
|
|
// 4. unclosed tag at end of response
|
|
text4 := "I will query:\n<tool_call>\n{\"name\": \"fetch_data\", \"id\": 123}"
|
|
tcs4, cleaned4, ok4 := DetectToolCalls(text4)
|
|
if !ok4 || len(tcs4) != 1 || tcs4[0].Function.Name != "fetch_data" || cleaned4 != "I will query:" {
|
|
t.Fatalf("DetectToolCalls 4 unclosed failed: ok=%v, tcs=%+v, cleaned=%q", ok4, tcs4, cleaned4)
|
|
}
|
|
|
|
// 5. function call syntax: get_weather(location="Tokyo")
|
|
text5 := "<tool_call>get_weather(location=\"Tokyo\")</tool_call>"
|
|
tcs5, _, ok5 := DetectToolCalls(text5)
|
|
if !ok5 || len(tcs5) != 1 || tcs5[0].Function.Name != "get_weather" || !strings.Contains(tcs5[0].Function.Arguments, "Tokyo") {
|
|
t.Fatalf("DetectToolCalls 5 function syntax failed: ok=%v, tcs=%+v", ok5, tcs5)
|
|
}
|
|
|
|
// 6. markdown ```json with tool payload
|
|
text6 := "```json\n{\"name\": \"lookup\", \"arguments\": {\"id\": 42}}\n```"
|
|
tcs6, _, ok6 := DetectToolCalls(text6)
|
|
if !ok6 || len(tcs6) != 1 || tcs6[0].Function.Name != "lookup" {
|
|
t.Fatalf("DetectToolCalls 6 json block failed: ok=%v, tcs=%+v", ok6, tcs6)
|
|
}
|
|
}
|
|
|
|
func TestStreamToolInterceptor(t *testing.T) {
|
|
w := httptest.NewRecorder()
|
|
streamer := NewStreamer(w, nil, "test-id", 123456789, "kimi-k3")
|
|
interceptor := NewStreamToolInterceptor(streamer)
|
|
|
|
interceptor.ProcessContentDelta("Checking weather for you...\n<tool_")
|
|
interceptor.ProcessContentDelta("call>\n{\"name\": \"get_weather\", \"arguments\": {\"city\": \"SF\"}}\n</tool_call>\nHave a nice day!")
|
|
interceptor.FlushRemaining()
|
|
|
|
if !interceptor.HasToolCalls() {
|
|
t.Fatalf("expected tool call registered in interceptor")
|
|
}
|
|
body := w.Body.String()
|
|
if !strings.Contains(body, "Checking weather for you...") || !strings.Contains(body, "get_weather") || strings.Contains(body, "<tool_call>") {
|
|
t.Fatalf("unexpected stream body: %s", body)
|
|
}
|
|
}
|
|
|
|
func TestEndToEndMockChatCompletion(t *testing.T) {
|
|
mockGradio := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path == "/gradio_api/call/on_submit" && r.Method == "POST" {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(GradioJoinResponse{EventID: "evt_12345"})
|
|
return
|
|
}
|
|
if r.URL.Path == "/gradio_api/call/on_submit/evt_12345" && r.Method == "GET" {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
flusher, _ := w.(http.Flusher)
|
|
fmt.Fprintf(w, "event: generating\ndata: [[{\"role\":\"user\",\"metadata\":null,\"content\":[{\"text\":\"Hi\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\"<\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":{\"title\":\"Reasoning\",\"status\":\"done\"},\"content\":[{\"text\":\">Thinking\"}],\"options\":null}]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
fmt.Fprintf(w, "event: complete\ndata: [[{\"role\":\"user\",\"metadata\":null,\"content\":[{\"text\":\"Hi\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\"<\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":{\"title\":\"Reasoning\",\"status\":\"done\"},\"content\":[{\"text\":\">Thinking</\"}],\"options\":null},{\"role\":\"assistant\",\"metadata\":null,\"content\":[{\"text\":\">\\n\\nHello!\"}],\"options\":null}]]\n\n")
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return
|
|
}
|
|
http.NotFound(w, r)
|
|
}))
|
|
defer mockGradio.Close()
|
|
|
|
svc := NewKimiService(mockGradio.URL, "kimi-k3", "direct:together", "max", "", 8192, 0.7)
|
|
mux := NewMux(svc)
|
|
|
|
// 1. Non-Streaming
|
|
reqB, _ := json.Marshal(ChatCompletionRequest{Model: "kimi-k3", Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, Stream: false})
|
|
req := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(reqB))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
w := httptest.NewRecorder()
|
|
mux.ServeHTTP(w, req)
|
|
|
|
var resp ChatCompletionResponse
|
|
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil || len(resp.Choices) == 0 {
|
|
t.Fatalf("non-streaming chat failed: %v", err)
|
|
}
|
|
if resp.Choices[0].Message.ReasoningContent != "Thinking" || resp.Choices[0].Message.GetContentString() != "Hello!" {
|
|
t.Errorf("unexpected choice response: %+v", resp.Choices[0])
|
|
}
|
|
|
|
// 2. Streaming
|
|
reqBStream, _ := json.Marshal(ChatCompletionRequest{Model: "kimi-k3", Messages: []ChatMessage{{Role: "user", Content: "Hi"}}, Stream: true})
|
|
reqStream := httptest.NewRequest("POST", "/v1/chat/completions", bytes.NewBuffer(reqBStream))
|
|
reqStream.Header.Set("Content-Type", "application/json")
|
|
wStream := httptest.NewRecorder()
|
|
mux.ServeHTTP(wStream, reqStream)
|
|
|
|
scanner := bufio.NewScanner(wStream.Body)
|
|
hasReasoning, hasContent, hasDone := false, false, false
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if strings.HasPrefix(line, "data: ") {
|
|
data := strings.TrimPrefix(line, "data: ")
|
|
if data == "[DONE]" {
|
|
hasDone = true
|
|
continue
|
|
}
|
|
var chunk StreamResponse
|
|
if err := json.Unmarshal([]byte(data), &chunk); err == nil && len(chunk.Choices) > 0 {
|
|
if chunk.Choices[0].Delta.ReasoningContent != "" {
|
|
hasReasoning = true
|
|
}
|
|
if chunk.Choices[0].Delta.Content != "" {
|
|
hasContent = true
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if !hasReasoning || !hasContent || !hasDone {
|
|
t.Errorf("streaming test failed: reasoning=%v, content=%v, done=%v", hasReasoning, hasContent, hasDone)
|
|
}
|
|
}
|