File size: 5,777 Bytes
2986042 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | package app
import (
"context"
"io"
"net/http"
"testing"
"time"
"ccLoad/internal/model"
"ccLoad/internal/testutil"
)
func TestTestChannelAPI_StreamIncludesUsageAndCost(t *testing.T) {
upstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/v1/messages" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
time.Sleep(20 * time.Millisecond)
// 模拟Claude风格SSE：usage在message_start/message_delta给出，内容在content_block_delta给出
_, _ = io.WriteString(w, "event: message_start\n")
_, _ = io.WriteString(w, "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10,\"output_tokens\":0,\"cache_read_input_tokens\":5,\"cache_creation\":{\"ephemeral_5m_input_tokens\":3,\"ephemeral_1h_input_tokens\":2}}}}\n\n")
_, _ = io.WriteString(w, "event: content_block_delta\n")
_, _ = io.WriteString(w, "data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"hi\"}}\n\n")
_, _ = io.WriteString(w, "event: message_delta\n")
_, _ = io.WriteString(w, "data: {\"type\":\"message_delta\",\"usage\":{\"input_tokens\":10,\"output_tokens\":20,\"cache_read_input_tokens\":5,\"cache_creation_input_tokens\":5,\"cache_creation\":{\"ephemeral_5m_input_tokens\":3,\"ephemeral_1h_input_tokens\":2}}}\n\n")
time.Sleep(20 * time.Millisecond)
_, _ = io.WriteString(w, "event: message_stop\n")
_, _ = io.WriteString(w, "data: {\"type\":\"message_stop\"}\n\n")
}))
defer upstream.Close()
srv := newInMemoryServer(t)
srv.client = upstream.Client()
cfg := &model.Config{
ID: 1,
Name: "test-channel",
URL: upstream.URL,
Priority: 1,
ModelEntries: []model.ModelEntry{{Model: "claude-3-haiku", RedirectModel: ""}},
ChannelType: "anthropic",
Enabled: true,
}
req := &testutil.TestChannelRequest{
Model: "claude-3-haiku",
Stream: true,
Content: "hi",
ChannelType: "anthropic",
}
result := srv.testChannelAPI(context.Background(), cfg, "sk-test", req)
if success, _ := result["success"].(bool); !success {
t.Fatalf("expected success, got: %#v", result)
}
if result["response_text"] != "hi" {
t.Fatalf("expected response_text=hi, got: %#v", result["response_text"])
}
apiResp, ok := result["api_response"].(map[string]any)
if !ok || apiResp == nil {
t.Fatalf("expected api_response, got: %#v", result["api_response"])
}
usage, ok := apiResp["usage"].(map[string]any)
if !ok || usage == nil {
t.Fatalf("expected api_response.usage, got: %#v", apiResp["usage"])
}
if usage["input_tokens"] == nil || usage["output_tokens"] == nil {
t.Fatalf("expected usage tokens, got: %#v", usage)
}
cost, ok := result["cost_usd"].(float64)
if !ok {
t.Fatalf("expected cost_usd(float64), got: %#v", result["cost_usd"])
}
if cost <= 0 {
t.Fatalf("expected cost_usd > 0, got: %v", cost)
}
firstByteDurationMs, ok := result["first_byte_duration_ms"].(int64)
if !ok || firstByteDurationMs <= 0 {
t.Fatalf("expected first_byte_duration_ms(int64)>0, got: %#v", result["first_byte_duration_ms"])
}
totalDurationMs, ok := result["duration_ms"].(int64)
if !ok || totalDurationMs <= 0 {
t.Fatalf("expected duration_ms(int64)>0, got: %#v", result["duration_ms"])
}
if totalDurationMs < firstByteDurationMs {
t.Fatalf("expected duration_ms>=first_byte_duration_ms, got %d < %d", totalDurationMs, firstByteDurationMs)
}
}
func TestTestChannelAPI_GeminiStreamIncludesTTFBAndText(t *testing.T) {
upstream := newTestHTTPServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Gemini 流式端点: /v1beta/models/{model}:streamGenerateContent
if r.URL.Path != "/v1beta/models/gemini-2.5-flash-lite:streamGenerateContent" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
time.Sleep(20 * time.Millisecond)
// Gemini SSE: candidates[0].content.parts[0].text, usage在usageMetadata中
_, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Hello\"}],\"role\":\"model\"}}],\"modelVersion\":\"gemini-2.5-flash-lite\"}\n\n")
_, _ = io.WriteString(w, "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\" world\"}],\"role\":\"model\"},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":10,\"candidatesTokenCount\":20,\"totalTokenCount\":30},\"modelVersion\":\"gemini-2.5-flash-lite\"}\n\n")
}))
defer upstream.Close()
srv := newInMemoryServer(t)
srv.client = upstream.Client()
cfg := &model.Config{
ID: 1,
Name: "gemini-channel",
URL: upstream.URL,
Priority: 1,
ModelEntries: []model.ModelEntry{{Model: "gemini-2.5-flash-lite"}},
ChannelType: "gemini",
Enabled: true,
}
req := &testutil.TestChannelRequest{
Model: "gemini-2.5-flash-lite",
Stream: true,
Content: "hi",
ChannelType: "gemini",
}
result := srv.testChannelAPI(context.Background(), cfg, "test-key", req)
if success, _ := result["success"].(bool); !success {
t.Fatalf("expected success, got: %#v", result)
}
// 验证文本提取
if result["response_text"] != "Hello world" {
t.Fatalf("expected response_text='Hello world', got: %#v", result["response_text"])
}
// 验证 TTFB
firstByteDurationMs, ok := result["first_byte_duration_ms"].(int64)
if !ok || firstByteDurationMs <= 0 {
t.Fatalf("expected first_byte_duration_ms(int64)>0, got: %#v", result["first_byte_duration_ms"])
}
// 验证总耗时
totalDurationMs, ok := result["duration_ms"].(int64)
if !ok || totalDurationMs <= 0 {
t.Fatalf("expected duration_ms(int64)>0, got: %#v", result["duration_ms"])
}
}
|