Files

770 lines
25 KiB
Go
Raw Permalink Normal View History

// qflash test suite
// Created by Luxferre in 2026, released into the public domain
package main
import (
"bufio"
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestChatMessageGetContentString(t *testing.T) {
// String content
msg1 := ChatMessage{Role: "user", Content: "Hello world"}
if msg1.GetContentString() != "Hello world" {
t.Fatalf("expected 'Hello world', got %q", msg1.GetContentString())
}
// Multi-part content
msg2 := ChatMessage{
Role: "user",
Content: []interface{}{
map[string]interface{}{"type": "text", "text": "Part 1 "},
map[string]interface{}{"type": "text", "text": "Part 2"},
},
}
if msg2.GetContentString() != "Part 1 Part 2" {
t.Fatalf("expected 'Part 1 Part 2', got %q", msg2.GetContentString())
}
// Nil content
msg3 := ChatMessage{Role: "assistant", Content: nil}
if msg3.GetContentString() != "" {
t.Fatalf("expected empty string, got %q", msg3.GetContentString())
}
}
func TestSOCKS5Parsing(t *testing.T) {
cfg, err := ParseSOCKS5URL("socks5://user:pass@127.0.0.1:9050")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg.Address != "127.0.0.1:9050" || cfg.Username != "user" || cfg.Password != "pass" {
t.Fatalf("mismatched parsed socks5 config: %+v", cfg)
}
cfg2, err := ParseSOCKS5URL("socks5h://proxy.internal:1080")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg2.Address != "proxy.internal:1080" || cfg2.Username != "" {
t.Fatalf("mismatched parsed socks5 config: %+v", cfg2)
}
cfg3, err := ParseSOCKS5URL("10.0.0.5")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if cfg3.Address != "10.0.0.5:1080" {
t.Fatalf("expected default port 1080, got %s", cfg3.Address)
}
}
func TestEffectiveModelID(t *testing.T) {
def := "Qwen/Qwen3.8-27B"
cases := map[string]string{
"": def,
"qwen": def,
"qwen3.8-flash-next": "Qwen/Qwen3.8-Flash-Next",
"qwen-flash-next": "Qwen/Qwen3.8-Flash-Next",
"qwen-flash": "Qwen/Qwen3.8-Flash-Next",
"Qwen/Qwen3.8-Flash-Next": "Qwen/Qwen3.8-Flash-Next",
"qwen3.8-27b": "Qwen/Qwen3.8-27B",
"qwen-27b": "Qwen/Qwen3.8-27B",
"Qwen/Qwen3.8-27B": "Qwen/Qwen3.8-27B",
"qwen3.8-27b-uncensored": "Qwen/Qwen3.8-27B-Uncensored",
"Qwen/Qwen3.8-27B-Uncensored": "Qwen/Qwen3.8-27B-Uncensored",
"custom-org/my-model": "custom-org/my-model",
}
for in, exp := range cases {
res := EffectiveModelID(in, def)
if res != exp {
t.Errorf("EffectiveModelID(%q) = %q; expected %q", in, res, exp)
}
}
}
func TestSeparateReasoningAndContentThinkTags(t *testing.T) {
raw := "<think>\nAnalyzing prompt step by step.\n</think>\n\nHere is the answer."
reasoning, content := SeparateReasoningAndContent(raw)
if reasoning != "Analyzing prompt step by step." {
t.Fatalf("unexpected reasoning: %q", reasoning)
}
if content != "Here is the answer." {
t.Fatalf("unexpected content: %q", content)
}
}
func TestSeparateReasoningAndContentGradioFormat(t *testing.T) {
raw := "> 💭 **Thinking Process (QSA Micro-block Reasoning):**\n>\n> Thinking Process:\n>\n> 1. Step one\n> 2. Step two\n\n---\n\n### Answer Header\n\nDetailed answer here."
reasoning, content := SeparateReasoningAndContent(raw)
if !strings.Contains(reasoning, "1. Step one") || !strings.Contains(reasoning, "2. Step two") {
t.Fatalf("expected reasoning to contain steps, got: %q", reasoning)
}
if strings.Contains(reasoning, ">") {
t.Fatalf("expected blockquote markers to be stripped, got: %q", reasoning)
}
if content != "### Answer Header\n\nDetailed answer here." {
t.Fatalf("unexpected content: %q", content)
}
}
func TestSeparateReasoningAndContentStreamingDivider(t *testing.T) {
raw := "> 💭 **Thinking Process (QSA Micro-block Reasoning):**\n>\n> Thinking Process:\n>\n> 1. Formulating response...\n\n---\n*Generating response...*"
reasoning, content := SeparateReasoningAndContent(raw)
if !strings.Contains(reasoning, "1. Formulating response...") {
t.Fatalf("expected reasoning, got: %q", reasoning)
}
if content != "" {
t.Fatalf("expected empty content during thought phase, got: %q", content)
}
}
func TestToolCallParsingAndDetection(t *testing.T) {
rawJSON := `{"name": "get_weather", "arguments": {"city": "Tokyo"}}`
tc, ok := parseSingleToolCall(rawJSON)
if !ok {
t.Fatalf("expected successful single tool call parse")
}
if tc.Function.Name != "get_weather" {
t.Fatalf("expected 'get_weather', got %q", tc.Function.Name)
}
rawXML := `<tool_call>
{"name": "fetch_data", "arguments": "{\"id\": 42}"}
</tool_call>`
calls, rem, hasCalls := DetectToolCalls(rawXML)
if !hasCalls || len(calls) != 1 {
t.Fatalf("expected 1 detected tool call, got %d", len(calls))
}
if calls[0].Function.Name != "fetch_data" {
t.Fatalf("expected 'fetch_data', got %q", calls[0].Function.Name)
}
if rem != "" {
t.Fatalf("expected empty remaining content, got %q", rem)
}
}
func TestStreamToolCallFilterNoLeak(t *testing.T) {
filter := &StreamToolCallFilter{}
var streamedContent strings.Builder
var emittedCalls []ToolCall
onContent := func(s string) {
streamedContent.WriteString(s)
}
onTool := func(tc ToolCall) {
emittedCalls = append(emittedCalls, tc)
}
// Stream in small split chunks that split the <tool_call> tag
chunks := []string{
"Here is the data: ",
"<tool",
"_call>\n",
`{"name": "query_db", "arguments": {"sql": "SELECT 1"}}`,
"\n</tool",
"_call>",
}
for _, c := range chunks {
filter.Feed(c, onContent, onTool)
}
filter.Flush(onContent, onTool)
if strings.Contains(streamedContent.String(), "<tool_call>") || strings.Contains(streamedContent.String(), "</tool_call>") {
t.Fatalf("tool call tags leaked into content: %q", streamedContent.String())
}
if streamedContent.String() != "Here is the data: " {
t.Fatalf("unexpected content: %q", streamedContent.String())
}
if len(emittedCalls) != 1 {
t.Fatalf("expected 1 emitted tool call, got %d", len(emittedCalls))
}
if emittedCalls[0].Function.Name != "query_db" {
t.Fatalf("expected 'query_db', got %q", emittedCalls[0].Function.Name)
}
}
func TestParseAssistantText(t *testing.T) {
dataJSON := `[[
{"role": "user", "metadata": null, "content": [{"text": "hi", "type": "text"}], "options": null},
{"role": "assistant", "metadata": null, "content": [{"text": "Hello, human!", "type": "text"}], "options": null}
]]`
text, ok := parseAssistantText(dataJSON)
if !ok {
t.Fatalf("expected successful parse of assistant text")
}
if text != "Hello, human!" {
t.Fatalf("expected 'Hello, human!', got %q", text)
}
}
func TestParseAssistantTextMultiMessageReasoning(t *testing.T) {
dataJSON := `[[
{"role": "user", "metadata": null, "content": [{"text": "Say hi", "type": "text"}]},
{"role": "assistant", "metadata": {"title": "Reasoning"}, "content": [{"text": "First, consider greeting politely.", "type": "text"}]},
{"role": "assistant", "metadata": null, "content": [{"text": "Hi there!", "type": "text"}]}
]]`
text, ok := parseAssistantText(dataJSON)
if !ok {
t.Fatalf("expected successful parse of multi-message assistant text")
}
reasoning, content := SeparateReasoningAndContent(text)
if reasoning != "First, consider greeting politely." {
t.Fatalf("expected reasoning 'First, consider greeting politely.', got %q", reasoning)
}
if content != "Hi there!" {
t.Fatalf("expected content 'Hi there!', got %q", content)
}
}
func TestQwenServiceChatMock(t *testing.T) {
// Mock upstream Gradio space server
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/gradio_api/call/chat_response" && r.Method == http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"event_id": "test_event_123"}`))
return
}
if r.URL.Path == "/gradio_api/call/chat_response/test_event_123" && r.Method == http.MethodGet {
w.Header().Set("Content-Type", "text/event-stream")
flusher, _ := w.(http.Flusher)
// Step 1: Thinking progress
chunk1 := `event: generating` + "\n" +
`data: [[{"role": "user", "content": [{"text": "Hello", "type": "text"}]}, {"role": "assistant", "content": [{"text": "> 💭 **Thinking Process:**\n>\n> Thinking Process:\n>\n> 1. Step 1\n\n---\n*Generating response...*", "type": "text"}]}]]` + "\n\n"
w.Write([]byte(chunk1))
flusher.Flush()
// Step 2: Final completion
chunk2 := `event: complete` + "\n" +
`data: [[{"role": "user", "content": [{"text": "Hello", "type": "text"}]}, {"role": "assistant", "content": [{"text": "> 💭 **Thinking Process:**\n>\n> Thinking Process:\n>\n> 1. Step 1\n\n---\n\nGreetings from mock Qwen!", "type": "text"}]}]]` + "\n\n"
w.Write([]byte(chunk2))
flusher.Flush()
return
}
http.NotFound(w, r)
}))
defer mockServer.Close()
svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-Flash-Next", "chat_response", "", "", "", "", true, true)
// 1. Test Non-streaming completion
rec := httptest.NewRecorder()
req := ChatCompletionRequest{
Model: "qwen3.8-flash-next",
Messages: []ChatMessage{
{Role: "user", Content: "Hello"},
},
Stream: false,
}
err := svc.Chat(rec, nil, req)
if err != nil {
t.Fatalf("unexpected error in Chat non-streaming: %v", err)
}
if rec.Code != http.StatusOK {
t.Fatalf("expected HTTP 200, got %d", rec.Code)
}
var resp ChatCompletionResponse
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
t.Fatalf("failed to decode completion response: %v", err)
}
if len(resp.Choices) == 0 {
t.Fatalf("expected choices, got 0")
}
if resp.Choices[0].Message.Content != "Greetings from mock Qwen!" {
t.Fatalf("unexpected message content: %v", resp.Choices[0].Message.Content)
}
if !strings.Contains(resp.Choices[0].Message.ReasoningContent, "1. Step 1") {
t.Fatalf("unexpected reasoning content: %v", resp.Choices[0].Message.ReasoningContent)
}
// 2. Test Streaming completion
recStream := httptest.NewRecorder()
reqStream := ChatCompletionRequest{
Model: "qwen-flash",
Messages: []ChatMessage{
{Role: "user", Content: "Hello"},
},
Stream: true,
}
errStream := svc.Chat(recStream, nil, reqStream)
if errStream != nil {
t.Fatalf("unexpected error in Chat streaming: %v", errStream)
}
scanner := bufio.NewScanner(recStream.Body)
var receivedReasoning strings.Builder
var receivedContent strings.Builder
var sawDone bool
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "data: ") {
payload := strings.TrimPrefix(line, "data: ")
if payload == "[DONE]" {
sawDone = true
continue
}
var sResp StreamResponse
if err := json.Unmarshal([]byte(payload), &sResp); err == nil && len(sResp.Choices) > 0 {
delta := sResp.Choices[0].Delta
if delta.ReasoningContent != "" {
receivedReasoning.WriteString(delta.ReasoningContent)
}
if delta.Content != "" {
receivedContent.WriteString(delta.Content)
}
}
}
}
if !sawDone {
t.Fatalf("expected [DONE] chunk in stream")
}
if !strings.Contains(receivedReasoning.String(), "1. Step 1") {
t.Fatalf("expected streamed reasoning, got %q", receivedReasoning.String())
}
if !strings.Contains(receivedContent.String(), "Greetings from mock Qwen!") {
t.Fatalf("expected streamed content, got %q", receivedContent.String())
}
}
func TestParseAssistantTextDirectString(t *testing.T) {
dataJSON := `["Step-by-step reasoning\n</think>\n\nFinal answer here", null]`
text, ok := parseAssistantText(dataJSON)
if !ok {
t.Fatalf("expected successful parse of direct string assistant text")
}
if text != "Step-by-step reasoning\n</think>\n\nFinal answer here" {
t.Fatalf("unexpected text: %q", text)
}
reasoning, content := SeparateReasoningAndContent(text)
if reasoning != "Step-by-step reasoning" {
t.Fatalf("unexpected reasoning: %q", reasoning)
}
if content != "Final answer here" {
t.Fatalf("unexpected content: %q", content)
}
}
func TestQwenServiceChatRespondMock(t *testing.T) {
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/gradio_api/call/respond" && r.Method == http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"event_id": "respond_event_456"}`))
return
}
if r.URL.Path == "/gradio_api/call/respond/respond_event_456" && r.Method == http.MethodGet {
w.Header().Set("Content-Type", "text/event-stream")
flusher, _ := w.(http.Flusher)
chunk := `event: complete` + "\n" +
`data: ["Analyzing the problem.\n</think>\n\nThe solution is 42.", null]` + "\n\n"
w.Write([]byte(chunk))
flusher.Flush()
return
}
http.NotFound(w, r)
}))
defer mockServer.Close()
svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "respond", "", "", "", "", true, true)
// Non-streaming test
rec := httptest.NewRecorder()
req := ChatCompletionRequest{
Model: "qwen3.8-27b",
Messages: []ChatMessage{
{Role: "user", Content: "What is the answer?"},
},
Stream: false,
}
err := svc.Chat(rec, nil, req)
if err != nil {
t.Fatalf("unexpected error in respond non-streaming: %v", err)
}
if rec.Code != http.StatusOK {
t.Fatalf("expected HTTP 200, got %d", rec.Code)
}
var resp ChatCompletionResponse
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
t.Fatalf("failed to decode response: %v", err)
}
if resp.Choices[0].Message.Content != "The solution is 42." {
t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
}
if resp.Choices[0].Message.ReasoningContent != "Analyzing the problem." {
t.Fatalf("unexpected reasoning: %v", resp.Choices[0].Message.ReasoningContent)
}
// Streaming test
recStream := httptest.NewRecorder()
reqStream := ChatCompletionRequest{
Model: "qwen3.8-27b",
Messages: []ChatMessage{
{Role: "user", Content: "What is the answer?"},
},
Stream: true,
}
errStream := svc.Chat(recStream, nil, reqStream)
if errStream != nil {
t.Fatalf("unexpected error in respond streaming: %v", errStream)
}
scanner := bufio.NewScanner(recStream.Body)
var receivedContent strings.Builder
var receivedReasoning strings.Builder
var sawDone bool
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "data: ") {
payload := strings.TrimPrefix(line, "data: ")
if payload == "[DONE]" {
sawDone = true
continue
}
var sResp StreamResponse
if err := json.Unmarshal([]byte(payload), &sResp); err == nil && len(sResp.Choices) > 0 {
delta := sResp.Choices[0].Delta
if delta.ReasoningContent != "" {
receivedReasoning.WriteString(delta.ReasoningContent)
}
if delta.Content != "" {
receivedContent.WriteString(delta.Content)
}
}
}
}
if !sawDone {
t.Fatalf("expected [DONE] in stream")
}
if receivedReasoning.String() != "Analyzing the problem." {
t.Fatalf("unexpected streamed reasoning: %q", receivedReasoning.String())
}
if receivedContent.String() != "The solution is 42." {
t.Fatalf("unexpected streamed content: %q", receivedContent.String())
}
}
func TestQwenServiceDirectOpenAIMock(t *testing.T) {
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/v1/chat/completions" && r.Method == http.MethodPost {
var req ChatCompletionRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, "bad request", http.StatusBadRequest)
return
}
if !req.Stream {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(ChatCompletionResponse{
ID: "chatcmpl-openai-mock",
Object: "chat.completion",
Created: 1234567890,
Model: req.Model,
Choices: []ChatCompletionResponseChoice{
{
Index: 0,
Message: ChatMessage{
Role: "assistant",
Content: "Hello from native OpenAI upstream!",
},
FinishReason: "stop",
},
},
})
return
}
// Stream
w.Header().Set("Content-Type", "text/event-stream")
flusher, _ := w.(http.Flusher)
chunk := `data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello from "}}]}
data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"native stream!"}}]}
data: [DONE]
`
w.Write([]byte(chunk))
if flusher != nil {
flusher.Flush()
}
return
}
http.NotFound(w, r)
}))
defer mockServer.Close()
svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "openai", "", "", "", "", true, true)
// Non-streaming
rec := httptest.NewRecorder()
req := ChatCompletionRequest{
Model: "qwen-openai",
Messages: []ChatMessage{
{Role: "user", Content: "Hello"},
},
Stream: false,
}
if err := svc.Chat(rec, nil, req); err != nil {
t.Fatalf("unexpected error in direct openai non-streaming: %v", err)
}
if rec.Code != http.StatusOK {
t.Fatalf("expected HTTP 200, got %d", rec.Code)
}
var resp ChatCompletionResponse
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
t.Fatalf("failed to decode response: %v", err)
}
if resp.Choices[0].Message.Content != "Hello from native OpenAI upstream!" {
t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
}
// Streaming
recStream := httptest.NewRecorder()
reqStream := ChatCompletionRequest{
Model: "qwen-openai",
Messages: []ChatMessage{
{Role: "user", Content: "Hello"},
},
Stream: true,
}
if err := svc.Chat(recStream, nil, reqStream); err != nil {
t.Fatalf("unexpected error in direct openai streaming: %v", err)
}
if !strings.Contains(recStream.Body.String(), "native stream!") {
t.Fatalf("expected streamed content, got %q", recStream.Body.String())
}
}
func TestQwenServiceAutoFailover(t *testing.T) {
// Server 1 simulates ZeroGPU quota exceeded (returns HTTP 429)
server1Hits := 0
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
server1Hits++
w.WriteHeader(http.StatusTooManyRequests)
w.Write([]byte(`{"error": "ZeroGPU runs limit reached: quota exceeded"}`))
}))
defer server1.Close()
// Server 2 is healthy and serves completions
server2Hits := 0
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
server2Hits++
if r.URL.Path == "/v1/chat/completions" {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(ChatCompletionResponse{
ID: "chatcmpl-failover-mock",
Object: "chat.completion",
Created: 1234567890,
Choices: []ChatCompletionResponseChoice{
{
Index: 0,
Message: ChatMessage{
Role: "assistant",
Content: "Response from fallback server!",
},
FinishReason: "stop",
},
},
})
return
}
http.NotFound(w, r)
}))
defer server2.Close()
endpoints := []string{server1.URL, server2.URL}
svc := NewQwenService(endpoints, "Qwen/Qwen3.8-27B", "openai", "", "", "", "", true, true)
rec := httptest.NewRecorder()
req := ChatCompletionRequest{
Model: "qwen3.8-27b",
Messages: []ChatMessage{
{Role: "user", Content: "Failover test"},
},
Stream: false,
}
err := svc.Chat(rec, nil, req)
if err != nil {
t.Fatalf("expected failover to succeed, got error: %v", err)
}
if rec.Code != http.StatusOK {
t.Fatalf("expected HTTP 200 after failover, got %d", rec.Code)
}
var resp ChatCompletionResponse
if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
t.Fatalf("failed to decode failover response: %v", err)
}
if resp.Choices[0].Message.Content != "Response from fallback server!" {
t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
}
if server1Hits != 1 {
t.Fatalf("expected server1 to be attempted once, got %d", server1Hits)
}
if server2Hits != 1 {
t.Fatalf("expected server2 to be called once on failover, got %d", server2Hits)
}
// Second request should skip server 1 because it entered cooldown
rec2 := httptest.NewRecorder()
err2 := svc.Chat(rec2, nil, req)
if err2 != nil {
t.Fatalf("expected second request to succeed directly on server2: %v", err2)
}
if server1Hits != 1 {
t.Fatalf("expected server1 to be skipped during cooldown, but was hit %d times", server1Hits)
}
if server2Hits != 2 {
t.Fatalf("expected server2 to receive the second request, got hits: %d", server2Hits)
}
}
func TestResolveMaxTokens(t *testing.T) {
// Default when empty
if v := ResolveMaxTokens(ChatCompletionRequest{}); v != 4096 {
t.Fatalf("expected default 4096, got %d", v)
}
// Explicit max_tokens
if v := ResolveMaxTokens(ChatCompletionRequest{MaxTokens: 2048}); v != 2048 {
t.Fatalf("expected 2048, got %d", v)
}
// max_completion_tokens precedence when max_tokens is 0
if v := ResolveMaxTokens(ChatCompletionRequest{MaxCompletionTokens: 1024}); v != 1024 {
t.Fatalf("expected 1024, got %d", v)
}
// Cap at 32768
if v := ResolveMaxTokens(ChatCompletionRequest{MaxTokens: 65536}); v != 32768 {
t.Fatalf("expected cap at 32768, got %d", v)
}
}
func TestQwenServiceSandboxFailover(t *testing.T) {
server1Hits := 0
server2Hits := 0
// Server 1: Returns Gradio chat_response sandbox simulation
server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
server1Hits++
if r.URL.Path == "/gradio_api/call/chat_response" {
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"event_id": "sandbox_event"}`))
return
}
if r.URL.Path == "/gradio_api/call/chat_response/sandbox_event" {
w.Header().Set("Content-Type", "text/event-stream")
flusher, _ := w.(http.Flusher)
chunk := "event: complete\n" +
"data: [[{\"role\": \"user\", \"content\": [{\"text\": \"hello\", \"type\": \"text\"}]}, {\"role\": \"assistant\", \"content\": [{\"text\": \"> 💭 **Thinking Process (QSA Micro-block Reasoning):**\\n\\n### Qwen3.8-Flash-Next Sandbox Response\\n\\nTo connect to a live inference engine, enter your endpoint credentials.\", \"type\": \"text\"}]}]]\n\n"
w.Write([]byte(chunk))
flusher.Flush()
return
}
http.NotFound(w, r)
}))
defer server1.Close()
// Server 2: Standard OpenAI server returning live completion
server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
server2Hits++
w.Header().Set("Content-Type", "application/json")
w.Write([]byte(`{"choices":[{"message":{"role":"assistant","content":"Live completion from fallback server"},"finish_reason":"stop"}]}`))
}))
defer server2.Close()
svc := NewQwenService([]string{server1.URL, server2.URL}, "Qwen/Qwen3.8-Flash-Next", "auto", "", "", "", "", true, true)
// Explicitly assign modes for mock URLs
svc.endpoints[0].Mode = "chat_response"
svc.endpoints[1].Mode = "openai"
// 1. Non-streaming failover test
rec := httptest.NewRecorder()
req := ChatCompletionRequest{
Model: "qwen",
Messages: []ChatMessage{{Role: "user", Content: "hello"}},
Stream: false,
}
err := svc.Chat(rec, nil, req)
if err != nil {
t.Fatalf("expected failover to succeed, got error: %v", err)
}
var res ChatCompletionResponse
if err := json.Unmarshal(rec.Body.Bytes(), &res); err != nil {
t.Fatalf("failed to decode response: %v", err)
}
if res.Choices[0].Message.Content != "Live completion from fallback server" {
t.Fatalf("unexpected content from failover: %v", res.Choices[0].Message.Content)
}
if server1Hits != 2 { // 1 for /gradio_api/call, 1 for /gradio_api/call/sandbox_event
t.Fatalf("expected server1 to be attempted, got %d hits", server1Hits)
}
if server2Hits != 1 {
t.Fatalf("expected server2 to be reached on failover, got %d hits", server2Hits)
}
// 2. Verify server1 is in cooldown and next request goes directly to server2
rec2 := httptest.NewRecorder()
err2 := svc.Chat(rec2, nil, req)
if err2 != nil {
t.Fatalf("expected request during cooldown to succeed on server2: %v", err2)
}
if server2Hits != 2 {
t.Fatalf("expected server2 to receive the second request directly, got %d hits", server2Hits)
}
// 3. Test streaming failover when server1 is out of cooldown
svc.endpoints[0].CooldownUntil = time.Time{}
streamReq := ChatCompletionRequest{
Model: "qwen",
Messages: []ChatMessage{{Role: "user", Content: "hello"}},
Stream: true,
}
recStream := httptest.NewRecorder()
errStream := svc.Chat(recStream, nil, streamReq)
if errStream != nil {
t.Fatalf("expected streaming failover to succeed, got error: %v", errStream)
}
if server2Hits != 3 {
t.Fatalf("expected server2 to receive the streaming failover, got %d hits", server2Hits)
}
}