Cyrene_v2/internal/capability/mcp/transport_test.go

package mcp

import (
	"net/http"
	"net/http/httptest"
	"testing"

	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
)

// TestNewTransport_SelectsByType 覆盖显式传输类型与历史推断。
func TestNewTransport_SelectsByType(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name string
		cfg  MCPConfig
		want any
	}{
		{name: "http", cfg: MCPConfig{URL: "http://x", TransportType: "http"}, want: &mcpsdk.StreamableClientTransport{}},
		{name: "streamable 别名", cfg: MCPConfig{URL: "http://x", TransportType: "streamable"}, want: &mcpsdk.StreamableClientTransport{}},
		{name: "sse", cfg: MCPConfig{URL: "http://x", TransportType: "sse"}, want: &mcpsdk.SSEClientTransport{}},
		{name: "stdio", cfg: MCPConfig{Command: "x", TransportType: "stdio"}, want: &mcpsdk.CommandTransport{}},
		{name: "推断 stdio", cfg: MCPConfig{Command: "x"}, want: &mcpsdk.CommandTransport{}},
		{name: "推断 sse", cfg: MCPConfig{URL: "http://x"}, want: &mcpsdk.SSEClientTransport{}},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			tr, err := newTransport(tt.cfg)
			if err != nil {
				t.Fatalf("newTransport: %v", err)
			}
			if got, want := typeName(tr), typeName(tt.want); got != want {
				t.Fatalf("传输类型 = %s, 期望 %s", got, want)
			}
		})
	}
}

// TestNewTransport_NoTransport 两个字段都空时必须报 ErrNoTransport。
func TestNewTransport_NoTransport(t *testing.T) {
	t.Parallel()
	if _, err := newTransport(MCPConfig{Name: "s"}); err == nil {
		t.Fatal("没有 Command/URL 却选出了传输")
	}
	// 显式类型但缺必要字段同样要报错（不静默构造坏传输）。
	if _, err := newTransport(MCPConfig{TransportType: "http"}); err == nil {
		t.Fatal("TransportType=http 但 URL 为空却没有报错")
	}
	if _, err := newTransport(MCPConfig{TransportType: "stdio"}); err == nil {
		t.Fatal("TransportType=stdio 但 Command 为空却没有报错")
	}
}

// TestBearerTransport_SetsAuthorization 凭据以 Authorization: Bearer 下发，
// 并且不污染原始请求。
func TestBearerTransport_SetsAuthorization(t *testing.T) {
	t.Parallel()
	got := make(chan string, 1)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		got <- r.Header.Get("Authorization")
		w.WriteHeader(http.StatusOK)
	}))
	defer ts.Close()

	client := httpClientFor(MCPConfig{APIKey: "sk-1234"})
	if client == nil {
		t.Fatal("配了 APIKey 却没有返回自定义 http client")
	}
	req, err := http.NewRequest(http.MethodGet, ts.URL, nil)
	if err != nil {
		t.Fatalf("NewRequest: %v", err)
	}
	resp, err := client.Do(req)
	if err != nil {
		t.Fatalf("Do: %v", err)
	}
	_ = resp.Body.Close()
	if auth := <-got; auth != "Bearer sk-1234" {
		t.Fatalf("Authorization = %q, 期望 Bearer sk-1234", auth)
	}
	if req.Header.Get("Authorization") != "" {
		t.Fatal("原始请求被就地改写了")
	}
	if httpClientFor(MCPConfig{}) != nil {
		t.Fatal("无 APIKey 时应返回 nil（用 SDK 默认客户端）")
	}
}

// TestEnvWith_OverridesAndPreserves stdio 子进程环境变量叠加语义。
func TestEnvWith_OverridesAndPreserves(t *testing.T) {
	t.Parallel()
	base := []string{"PATH=/usr/bin", "HOME=/root"}
	got := envWith(base, map[string]string{"HOME": "/tmp", "TOKEN": "abc"})
	if len(got) != 3 {
		t.Fatalf("len = %d, 期望 3（%v）", len(got), got)
	}
	seen := map[string]bool{}
	for _, kv := range got {
		seen[kv] = true
	}
	if !seen["PATH=/usr/bin"] || !seen["HOME=/tmp"] || !seen["TOKEN=abc"] {
		t.Fatalf("合并后的环境 = %v", got)
	}
	if len(envWith(base, nil)) != len(base) {
		t.Fatal("空 extra 改变了基础环境")
	}
}

// typeName 返回传输的类别名，避免在断言里依赖具体指针比较。
func typeName(v any) string {
	switch v.(type) {
	case *mcpsdk.StreamableClientTransport:
		return "streamable"
	case *mcpsdk.SSEClientTransport:
		return "sse"
	case *mcpsdk.CommandTransport:
		return "stdio"
	default:
		return "unknown"
	}
}