FranklinWH-Toolbox/service/agentmcp/http_test.go

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
}