package agentmcp import ( "net/http" "net/http/httptest" "reflect" "regexp" "strings" "testing" "go.uber.org/zap" ) var toolNameRE = regexp.MustCompile(`"name":"([a-z_]+)"`) // listTools 以 tools/list 请求 path,返回状态码与工具名集合。 func listTools(t *testing.T, h http.Handler, path, token string) (int, map[string]bool) { t.Helper() req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"tools/list"}`)) req.Host = "127.0.0.1:8091" req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") req.Header.Set("Authorization", "Bearer "+token) rec := httptest.NewRecorder() h.ServeHTTP(rec, req) names := map[string]bool{} for _, m := range toolNameRE.FindAllStringSubmatch(rec.Body.String(), -1) { names[m[1]] = true } return rec.Code, names } // 一个端口按路径区分工具集:各路径只暴露自己的工具,令牌各自校验;/mcp 兼容为只读工具集。 func TestHandlerRoutesSetsByPath(t *testing.T) { log := zap.NewNop() var endpoints []Endpoint tokens := map[string]string{SetReadOnly: "tok-ro", SetOps: "tok-ops"} for _, set := range AllSets { s, tasks, err := NewServer(set, log, 2) if err != nil { t.Fatal(err) } endpoints = append(endpoints, Endpoint{Set: set, Server: s, Tasks: tasks, Token: tokens[set]}) } h := Handler(endpoints, log) roCount := len(mustTools(t, SetReadOnly)) opsCount := len(mustTools(t, SetOps)) for _, path := range []string{"/mcp", "/mcp/readonly"} { code, names := listTools(t, h, path, "tok-ro") if code != http.StatusOK || len(names) != roCount || !names["site_search"] || names["site_info_refresh"] { t.Errorf("%s: code=%d 工具数=%d,期望 200 与只读工具集 %d 个", path, code, len(names), roCount) } } code, names := listTools(t, h, "/mcp/ops", "tok-ops") if code != http.StatusOK || len(names) != opsCount || !names["site_info_refresh"] || names["site_search"] { t.Errorf("/mcp/ops: code=%d 工具数=%d,期望 200 与写工具集 %d 个", code, len(names), opsCount) } if code, _ := listTools(t, h, "/mcp/ops", "tok-ro"); code != http.StatusUnauthorized { t.Errorf("/mcp/ops 用只读令牌应返回 401,实际 %d", code) } if code, _ := listTools(t, h, "/mcp/readonly", "tok-ops"); code != http.StatusUnauthorized { t.Errorf("/mcp/readonly 用写令牌应返回 401,实际 %d", code) } if code, _ := listTools(t, h, "/mcp/unknown", "tok-ro"); code != http.StatusNotFound { t.Errorf("/mcp/unknown 应返回 404,实际 %d", code) } } // 只加载只读工具集时不暴露 /mcp/ops。 func TestHandlerOnlyMountsLoadedSets(t *testing.T) { log := zap.NewNop() s, tasks, err := NewServer(SetReadOnly, log, 2) if err != nil { t.Fatal(err) } h := Handler([]Endpoint{{Set: SetReadOnly, Server: s, Tasks: tasks, Token: "tok"}}, log) if code, _ := listTools(t, h, "/mcp/ops", "tok"); code != http.StatusNotFound { t.Errorf("未加载 ops 时 /mcp/ops 应返回 404,实际 %d", code) } } func TestParseSets(t *testing.T) { cases := []struct { in string want []string wantErr bool }{ {"", AllSets, false}, {"readonly", []string{SetReadOnly}, false}, {" ops , readonly ,ops", []string{SetOps, SetReadOnly}, false}, {"readonly,foo", nil, true}, {" , ", nil, true}, } for _, c := range cases { got, err := ParseSets(c.in) if (err != nil) != c.wantErr || (!c.wantErr && !reflect.DeepEqual(got, c.want)) { t.Errorf("ParseSets(%q) = %v, %v;期望 %v, err=%v", c.in, got, err, c.want, c.wantErr) } } } func TestTokenFor(t *testing.T) { s := Settings{Token: "base", SetTokens: map[string]string{SetOps: "ops-only"}} if s.TokenFor(SetReadOnly) != "base" || s.TokenFor(SetOps) != "ops-only" { t.Errorf("TokenFor 结果不符:readonly=%s ops=%s", s.TokenFor(SetReadOnly), s.TokenFor(SetOps)) } } func mustTools(t *testing.T, set string) []string { t.Helper() list, err := Tools(set) if err != nil { t.Fatal(err) } names := make([]string, 0, len(list)) for _, r := range list { names = append(names, r.Meta().Name) } return names }