|
|
@@ -0,0 +1,124 @@
|
|
|
+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
|
|
|
+}
|