hunjing/shenshi/internal/transform/adaptors_test.go

package transform

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"hunjing/shenshi/internal/ir"
)

func readJSON(r *http.Request, v any) error {
	defer func() { _ = r.Body.Close() }()
	return json.NewDecoder(r.Body).Decode(v)
}

// 收集全量增量（适配器测试通用）。
func collect(t *testing.T, ad Adaptor, req *ir.ChatRequest) []ir.Delta {
	t.Helper()
	out := make(chan ir.Delta, 64)
	done := make(chan error, 1)
	go func() {
		defer close(out)
		done <- ad.Stream(context.Background(), req, out)
	}()
	var deltas []ir.Delta
	for d := range out {
		deltas = append(deltas, d)
	}
	if err := <-done; err != nil {
		t.Fatalf("Stream 出错: %v", err)
	}
	return deltas
}

func TestAnthropicAdaptorRequestAndStream(t *testing.T) {
	var gotBody antRequest
	var gotAuth, gotVersion bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/messages" {
			t.Errorf("路径错误: %s", r.URL.Path)
		}
		gotAuth = r.Header.Get("x-api-key") == "sk-test"
		gotVersion = r.Header.Get("anthropic-version") != ""
		_ = readJSON(r, &gotBody)
		w.Header().Set("content-type", "text/event-stream")
		_, _ = w.Write([]byte("data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":11,\"output_tokens\":1}}}\n\n"))
		_, _ = w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"你好\"}}\n\n"))
		_, _ = w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"，镜中\"}}\n\n"))
		_, _ = w.Write([]byte("data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":7}}\n\n"))
	}))
	defer srv.Close()

	ad := &AnthropicAdaptor{EndpointName: "ant", BaseURL: srv.URL, APIKey: "sk-test", UpstreamModel: "claude-x"}
	req := &ir.ChatRequest{
		Messages: []ir.Message{
			{Role: "system", Content: "系统契约"},
			{Role: "user", Content: "照心"},
		},
		RunID: "run_ant", Stream: true,
	}
	deltas := collect(t, ad, req)

	if !gotAuth || !gotVersion {
		t.Fatal("缺少 x-api-key / anthropic-version 头")
	}
	if gotBody.System == nil || *gotBody.System != "系统契约" {
		t.Fatalf("system 应提升为顶层字段: %+v", gotBody.System)
	}
	if gotBody.MaxTokens == nil {
		t.Fatal("Anthropic 要求 max_tokens 必填，适配器应兜底")
	}
	var text string
	var last ir.Delta
	for _, d := range deltas {
		text += d.Delta
		if d.FinishReason != "" {
			last = d
		}
	}
	if text != "你好，镜中" {
		t.Fatalf("增量拼接错误: %q", text)
	}
	if last.FinishReason != "stop" || last.UsagePromptTokens != 11 || last.UsageCompletionTokens != 7 {
		t.Fatalf("末帧用量错误: %+v", last)
	}
}

func TestAnthropicAdaptorSurfaces429(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusTooManyRequests)
		_, _ = w.Write([]byte(`{"error":"rate_limited"}`))
	}))
	defer srv.Close()
	ad := &AnthropicAdaptor{BaseURL: srv.URL, APIKey: "k", UpstreamModel: "m"}
	out := make(chan ir.Delta, 4)
	err := ad.Stream(context.Background(), &ir.ChatRequest{RunID: "r", Stream: true, Messages: []ir.Message{{Role: "user", Content: "x"}}}, out)
	up, ok := err.(*UpstreamError)
	if !ok || up.StatusCode != 429 {
		t.Fatalf("应携带 429 UpstreamError 供绑定级冷却，实际 %v", err)
	}
}

func TestGeminiAdaptorRequestAndStream(t *testing.T) {
	var gotPath, gotKey string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		gotPath = r.URL.Path
		gotKey = r.Header.Get("x-goog-api-key")
		w.Header().Set("content-type", "text/event-stream")
		_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"第一\"}]}}]}\n\n"))
		_, _ = w.Write([]byte("data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"第二\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":5,\"candidatesTokenCount\":9}}\n\n"))
	}))
	defer srv.Close()

	ad := &GeminiAdaptor{EndpointName: "gem", BaseURL: srv.URL, APIKey: "AIza-test", UpstreamModel: "gemini-x"}
	req := &ir.ChatRequest{
		Messages: []ir.Message{
			{Role: "system", Content: "系统契约"},
			{Role: "user", Content: "问"},
		},
		RunID: "run_gem", Stream: true,
	}
	deltas := collect(t, ad, req)

	if gotPath != "/v1beta/models/gemini-x:streamGenerateContent" {
		t.Fatalf("Gemini 路径错误: %s", gotPath)
	}
	if gotKey != "AIza-test" {
		t.Fatal("缺少 x-goog-api-key 头")
	}
	var text string
	var last ir.Delta
	for _, d := range deltas {
		text += d.Delta
		if d.FinishReason != "" {
			last = d
		}
	}
	if text != "第一第二" {
		t.Fatalf("增量拼接错误: %q", text)
	}
	if last.UsagePromptTokens != 5 || last.UsageCompletionTokens != 9 {
		t.Fatalf("用量解析错误: %+v", last)
	}
}