150 lines
3.7 KiB
Go
150 lines
3.7 KiB
Go
package tools_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"testing"
|
|
|
|
"ollie/tools"
|
|
)
|
|
|
|
// stubServer is a minimal Server used to verify the contract.
|
|
type stubServer struct {
|
|
name string
|
|
tools []tools.ToolInfo
|
|
}
|
|
|
|
func (s *stubServer) ListTools() ([]tools.ToolInfo, error) { return s.tools, nil }
|
|
func (s *stubServer) CallTool(_ context.Context, tool string, _ json.RawMessage) (json.RawMessage, error) {
|
|
return json.RawMessage(`{"tool":"` + tool + `"}`), nil
|
|
}
|
|
|
|
func newStub(name string, toolNames ...string) *stubServer {
|
|
var ti []tools.ToolInfo
|
|
for _, n := range toolNames {
|
|
ti = append(ti, tools.ToolInfo{Name: n, Description: n + " desc"})
|
|
}
|
|
return &stubServer{name: name, tools: ti}
|
|
}
|
|
|
|
// checkServerContract verifies Server invariants.
|
|
func checkServerContract(t *testing.T, s tools.Server) {
|
|
t.Helper()
|
|
tl, err := s.ListTools()
|
|
if err != nil {
|
|
t.Fatalf("ListTools: %v", err)
|
|
}
|
|
if tl == nil {
|
|
t.Fatal("ListTools returned nil")
|
|
}
|
|
for _, ti := range tl {
|
|
res, err := s.CallTool(context.Background(), ti.Name, json.RawMessage(`{}`))
|
|
if err != nil {
|
|
t.Errorf("CallTool(%q): %v", ti.Name, err)
|
|
}
|
|
if len(res) == 0 {
|
|
t.Errorf("CallTool(%q) returned empty result", ti.Name)
|
|
}
|
|
}
|
|
}
|
|
|
|
// checkDispatcherContract verifies Dispatcher invariants.
|
|
func checkDispatcherContract(t *testing.T, d tools.Dispatcher, servers map[string]*stubServer) {
|
|
t.Helper()
|
|
|
|
for name, s := range servers {
|
|
d.AddServer(name, s)
|
|
}
|
|
|
|
// GetServer round-trip
|
|
for name := range servers {
|
|
s, ok := d.GetServer(name)
|
|
if !ok || s == nil {
|
|
t.Errorf("GetServer(%q) not found after AddServer", name)
|
|
}
|
|
}
|
|
if _, ok := d.GetServer("nonexistent"); ok {
|
|
t.Error("GetServer returned true for unregistered server")
|
|
}
|
|
|
|
// ListTools aggregates all servers
|
|
all, err := d.ListTools()
|
|
if err != nil {
|
|
t.Fatalf("ListTools: %v", err)
|
|
}
|
|
var wantCount int
|
|
for _, s := range servers {
|
|
wantCount += len(s.tools)
|
|
}
|
|
if len(all) != wantCount {
|
|
t.Errorf("ListTools returned %d tools, want %d", len(all), wantCount)
|
|
}
|
|
for _, ti := range all {
|
|
if ti.Server == "" {
|
|
t.Errorf("tool %q has empty Server field", ti.Name)
|
|
}
|
|
}
|
|
|
|
// Dispatch routes to correct server
|
|
for name, s := range servers {
|
|
for _, ti := range s.tools {
|
|
res, err := d.Dispatch(context.Background(), name, ti.Name, json.RawMessage(`{}`))
|
|
if err != nil {
|
|
t.Errorf("Dispatch(%q, %q): %v", name, ti.Name, err)
|
|
}
|
|
if len(res) == 0 {
|
|
t.Errorf("Dispatch(%q, %q) returned empty", name, ti.Name)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Dispatch to unknown server must error
|
|
_, err = d.Dispatch(context.Background(), "nonexistent", "tool", json.RawMessage(`{}`))
|
|
if err == nil {
|
|
t.Error("Dispatch to unknown server should error")
|
|
}
|
|
}
|
|
|
|
func TestStubServerContract(t *testing.T) {
|
|
checkServerContract(t, newStub("s", "a", "b"))
|
|
}
|
|
|
|
func TestDispatcherContract(t *testing.T) {
|
|
servers := map[string]*stubServer{
|
|
"alpha": newStub("alpha", "tool1", "tool2"),
|
|
"beta": newStub("beta", "tool3"),
|
|
}
|
|
checkDispatcherContract(t, tools.NewDispatcher(), servers)
|
|
}
|
|
|
|
// failServer is a Server whose ListTools always errors.
|
|
type failServer struct{ stubServer }
|
|
|
|
func (f *failServer) ListTools() ([]tools.ToolInfo, error) {
|
|
return nil, fmt.Errorf("boom")
|
|
}
|
|
|
|
func TestDispatcherListToolsError(t *testing.T) {
|
|
d := tools.NewDispatcher()
|
|
d.AddServer("bad", &failServer{})
|
|
_, err := d.ListTools()
|
|
if err == nil {
|
|
t.Error("expected error from failing server")
|
|
}
|
|
}
|
|
|
|
func TestNewDispatcherFunc(t *testing.T) {
|
|
factory := tools.NewDispatcherFunc(map[string]func() tools.Server{
|
|
"s1": func() tools.Server { return newStub("s1", "t1") },
|
|
})
|
|
d := factory()
|
|
tl, err := d.ListTools()
|
|
if err != nil {
|
|
t.Fatalf("ListTools: %v", err)
|
|
}
|
|
if len(tl) != 1 {
|
|
t.Errorf("got %d tools, want 1", len(tl))
|
|
}
|
|
}
|