feat: add go helper client and mcp tools
This commit is contained in:
zhaoyilun
235
internal/tools/winhelper_test.go
Normal file
235
internal/tools/winhelper_test.go
Normal file
@@ -0,0 +1,235 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"sort"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"isphere-ai-bridge/internal/helperclient"
|
||||
)
|
||||
|
||||
var expectedWinHelperToolNames = []string{
|
||||
"win_helper_version",
|
||||
"win_helper_self_check",
|
||||
"win_helper_scan_windows",
|
||||
"win_helper_dump_uia",
|
||||
}
|
||||
|
||||
func TestWinHelperToolsRegisterExactlyFourReadOnlyTools(t *testing.T) {
|
||||
fake := &fakeHelperCaller{}
|
||||
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
||||
RegisterWinHelperTools(server, fake)
|
||||
})
|
||||
defer cleanup()
|
||||
|
||||
result, err := session.ListTools(context.Background(), &mcp.ListToolsParams{})
|
||||
if err != nil {
|
||||
t.Fatalf("list tools: %v", err)
|
||||
}
|
||||
|
||||
gotNames := make([]string, 0, len(result.Tools))
|
||||
for _, tool := range result.Tools {
|
||||
gotNames = append(gotNames, tool.Name)
|
||||
lower := strings.ToLower(tool.Name)
|
||||
forbidden := []string{"send", "search", "file", "login", "conversation", "upload", "download", "receive"}
|
||||
for _, word := range forbidden {
|
||||
if strings.Contains(lower, word) {
|
||||
t.Fatalf("tool %q contains forbidden action word %q", tool.Name, word)
|
||||
}
|
||||
}
|
||||
}
|
||||
sort.Strings(gotNames)
|
||||
wantNames := append([]string(nil), expectedWinHelperToolNames...)
|
||||
sort.Strings(wantNames)
|
||||
if !reflect.DeepEqual(gotNames, wantNames) {
|
||||
t.Fatalf("registered tool names = %#v, want %#v", gotNames, wantNames)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWinHelperToolSchemasExposeOnlyAllowedParameters(t *testing.T) {
|
||||
assertJSONFields[VersionArgs](t)
|
||||
assertJSONFields[SelfCheckArgs](t)
|
||||
assertJSONFields[ScanWindowsArgs](t, "include_all_visible")
|
||||
assertJSONFields[DumpUIAArgs](t, "hwnd", "max_depth", "include_text", "max_children")
|
||||
|
||||
fake := &fakeHelperCaller{}
|
||||
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
||||
RegisterWinHelperTools(server, fake)
|
||||
})
|
||||
defer cleanup()
|
||||
|
||||
result, err := session.ListTools(context.Background(), &mcp.ListToolsParams{})
|
||||
if err != nil {
|
||||
t.Fatalf("list tools: %v", err)
|
||||
}
|
||||
schemas := map[string][]string{}
|
||||
for _, tool := range result.Tools {
|
||||
schemas[tool.Name] = schemaPropertyNames(t, tool.InputSchema)
|
||||
}
|
||||
|
||||
assertSameStrings(t, schemas["win_helper_version"], nil)
|
||||
assertSameStrings(t, schemas["win_helper_self_check"], nil)
|
||||
assertSameStrings(t, schemas["win_helper_scan_windows"], []string{"include_all_visible"})
|
||||
assertSameStrings(t, schemas["win_helper_dump_uia"], []string{"hwnd", "max_depth", "include_text", "max_children"})
|
||||
}
|
||||
|
||||
func TestWinHelperHandlersCallHelperClientAbstractionWithCorrectOps(t *testing.T) {
|
||||
fake := &fakeHelperCaller{}
|
||||
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
||||
RegisterWinHelperTools(server, fake)
|
||||
})
|
||||
defer cleanup()
|
||||
|
||||
calls := []struct {
|
||||
tool string
|
||||
args map[string]any
|
||||
}{
|
||||
{tool: "win_helper_version", args: map[string]any{}},
|
||||
{tool: "win_helper_self_check", args: map[string]any{}},
|
||||
{tool: "win_helper_scan_windows", args: map[string]any{"include_all_visible": true}},
|
||||
{tool: "win_helper_dump_uia", args: map[string]any{"hwnd": "0x001A0B2C", "max_depth": 3, "include_text": true, "max_children": 25}},
|
||||
}
|
||||
|
||||
for _, call := range calls {
|
||||
if _, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: call.tool, Arguments: call.args}); err != nil {
|
||||
t.Fatalf("call %s: %v", call.tool, err)
|
||||
}
|
||||
}
|
||||
|
||||
wantOps := []string{"version", "self_check", "scan_windows", "dump_uia"}
|
||||
if len(fake.calls) != len(wantOps) {
|
||||
t.Fatalf("helper calls = %#v, want %d calls", fake.calls, len(wantOps))
|
||||
}
|
||||
for i, wantOp := range wantOps {
|
||||
if fake.calls[i].op != wantOp {
|
||||
t.Fatalf("call %d op = %q, want %q", i, fake.calls[i].op, wantOp)
|
||||
}
|
||||
}
|
||||
assertSameMap(t, fake.calls[0].args, map[string]any{})
|
||||
assertSameMap(t, fake.calls[1].args, map[string]any{})
|
||||
assertSameMap(t, fake.calls[2].args, map[string]any{"include_all_visible": true})
|
||||
assertSameMap(t, fake.calls[3].args, map[string]any{"hwnd": "0x001A0B2C", "max_depth": float64(3), "include_text": true, "max_children": float64(25)})
|
||||
}
|
||||
|
||||
type fakeHelperCaller struct {
|
||||
calls []helperCall
|
||||
}
|
||||
|
||||
type helperCall struct {
|
||||
op string
|
||||
args map[string]any
|
||||
}
|
||||
|
||||
func (f *fakeHelperCaller) Call(_ context.Context, op string, args map[string]any) (*helperclient.HelperResponse, error) {
|
||||
copied := map[string]any{}
|
||||
for k, v := range args {
|
||||
copied[k] = v
|
||||
}
|
||||
f.calls = append(f.calls, helperCall{op: op, args: copied})
|
||||
return &helperclient.HelperResponse{
|
||||
Protocol: helperclient.HelperProtocolV1,
|
||||
RequestID: "fake-request",
|
||||
Op: op,
|
||||
OK: true,
|
||||
Data: map[string]json.RawMessage{
|
||||
"helper_name": json.RawMessage(`"ISphereWinHelper"`),
|
||||
},
|
||||
Warnings: []string{},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func connectToolsTestSession(t *testing.T, register func(*mcp.Server)) (*mcp.ClientSession, func()) {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
serverTransport, clientTransport := mcp.NewInMemoryTransports()
|
||||
server := mcp.NewServer(&mcp.Implementation{Name: "tools-test-server", Version: "0.0.0"}, nil)
|
||||
register(server)
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- server.Run(ctx, serverTransport)
|
||||
}()
|
||||
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "tools-test-client", Version: "0.0.0"}, nil)
|
||||
session, err := client.Connect(ctx, clientTransport, nil)
|
||||
if err != nil {
|
||||
cancel()
|
||||
t.Fatalf("connect client: %v", err)
|
||||
}
|
||||
|
||||
cleanup := func() {
|
||||
_ = session.Close()
|
||||
cancel()
|
||||
select {
|
||||
case <-errCh:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatalf("server did not stop")
|
||||
}
|
||||
}
|
||||
return session, cleanup
|
||||
}
|
||||
|
||||
func assertJSONFields[T any](t *testing.T, want ...string) {
|
||||
t.Helper()
|
||||
var zero T
|
||||
typeOf := reflect.TypeOf(zero)
|
||||
if typeOf.Kind() != reflect.Struct {
|
||||
t.Fatalf("%T is %s, want struct", zero, typeOf.Kind())
|
||||
}
|
||||
got := make([]string, 0, typeOf.NumField())
|
||||
for i := 0; i < typeOf.NumField(); i++ {
|
||||
field := typeOf.Field(i)
|
||||
name := strings.Split(field.Tag.Get("json"), ",")[0]
|
||||
if name == "" || name == "-" {
|
||||
continue
|
||||
}
|
||||
got = append(got, name)
|
||||
}
|
||||
assertSameStrings(t, got, want)
|
||||
}
|
||||
|
||||
func schemaPropertyNames(t *testing.T, schema any) []string {
|
||||
t.Helper()
|
||||
payload, err := json.Marshal(schema)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal schema: %v", err)
|
||||
}
|
||||
var decoded map[string]any
|
||||
if err := json.Unmarshal(payload, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal schema %s: %v", payload, err)
|
||||
}
|
||||
properties, _ := decoded["properties"].(map[string]any)
|
||||
if len(properties) == 0 {
|
||||
return nil
|
||||
}
|
||||
names := make([]string, 0, len(properties))
|
||||
for name := range properties {
|
||||
names = append(names, name)
|
||||
}
|
||||
sort.Strings(names)
|
||||
return names
|
||||
}
|
||||
|
||||
func assertSameStrings(t *testing.T, got []string, want []string) {
|
||||
t.Helper()
|
||||
gotCopy := append([]string(nil), got...)
|
||||
wantCopy := append([]string(nil), want...)
|
||||
sort.Strings(gotCopy)
|
||||
sort.Strings(wantCopy)
|
||||
if !reflect.DeepEqual(gotCopy, wantCopy) {
|
||||
t.Fatalf("strings = %#v, want %#v", gotCopy, wantCopy)
|
||||
}
|
||||
}
|
||||
|
||||
func assertSameMap(t *testing.T, got map[string]any, want map[string]any) {
|
||||
t.Helper()
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("map = %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user