ini upl

2026-09-02 16:19:51 +03:00
commit 11185339a9
6 changed files with 2202 additions and 0 deletions
@@ -0,0 +1,396 @@
+// 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)
+  }
+}