115 lines
2.9 KiB
Go
115 lines
2.9 KiB
Go
|
package mcp
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||
|
|
)
|
||
|
|
|
||
|
|
// stubSubsystem is a minimal Subsystem for testing.
|
||
|
|
type stubSubsystem struct {
|
||
|
|
name string
|
||
|
|
toolsRegistered bool
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *stubSubsystem) Name() string { return s.name }
|
||
|
|
|
||
|
|
func (s *stubSubsystem) RegisterTools(server *mcp.Server) {
|
||
|
|
s.toolsRegistered = true
|
||
|
|
}
|
||
|
|
|
||
|
|
// shutdownSubsystem tracks Shutdown calls.
|
||
|
|
type shutdownSubsystem struct {
|
||
|
|
stubSubsystem
|
||
|
|
shutdownCalled bool
|
||
|
|
shutdownErr error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *shutdownSubsystem) Shutdown(_ context.Context) error {
|
||
|
|
s.shutdownCalled = true
|
||
|
|
return s.shutdownErr
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWithSubsystem_Good_Registration(t *testing.T) {
|
||
|
|
sub := &stubSubsystem{name: "test-sub"}
|
||
|
|
svc, err := New(WithSubsystem(sub))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
if len(svc.Subsystems()) != 1 {
|
||
|
|
t.Fatalf("expected 1 subsystem, got %d", len(svc.Subsystems()))
|
||
|
|
}
|
||
|
|
if svc.Subsystems()[0].Name() != "test-sub" {
|
||
|
|
t.Errorf("expected name 'test-sub', got %q", svc.Subsystems()[0].Name())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWithSubsystem_Good_ToolsRegistered(t *testing.T) {
|
||
|
|
sub := &stubSubsystem{name: "tools-sub"}
|
||
|
|
_, err := New(WithSubsystem(sub))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
if !sub.toolsRegistered {
|
||
|
|
t.Error("expected RegisterTools to have been called")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWithSubsystem_Good_MultipleSubsystems(t *testing.T) {
|
||
|
|
sub1 := &stubSubsystem{name: "sub-1"}
|
||
|
|
sub2 := &stubSubsystem{name: "sub-2"}
|
||
|
|
svc, err := New(WithSubsystem(sub1), WithSubsystem(sub2))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
if len(svc.Subsystems()) != 2 {
|
||
|
|
t.Fatalf("expected 2 subsystems, got %d", len(svc.Subsystems()))
|
||
|
|
}
|
||
|
|
if !sub1.toolsRegistered || !sub2.toolsRegistered {
|
||
|
|
t.Error("expected all subsystems to have RegisterTools called")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSubsystemShutdown_Good(t *testing.T) {
|
||
|
|
sub := &shutdownSubsystem{stubSubsystem: stubSubsystem{name: "shutdown-sub"}}
|
||
|
|
svc, err := New(WithSubsystem(sub))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
if err := svc.Shutdown(context.Background()); err != nil {
|
||
|
|
t.Fatalf("Shutdown() failed: %v", err)
|
||
|
|
}
|
||
|
|
if !sub.shutdownCalled {
|
||
|
|
t.Error("expected Shutdown to have been called")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSubsystemShutdown_Bad_Error(t *testing.T) {
|
||
|
|
sub := &shutdownSubsystem{
|
||
|
|
stubSubsystem: stubSubsystem{name: "fail-sub"},
|
||
|
|
shutdownErr: context.DeadlineExceeded,
|
||
|
|
}
|
||
|
|
svc, err := New(WithSubsystem(sub))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
err = svc.Shutdown(context.Background())
|
||
|
|
if err == nil {
|
||
|
|
t.Fatal("expected error from Shutdown")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestSubsystemShutdown_Good_NoShutdownInterface(t *testing.T) {
|
||
|
|
// A plain Subsystem (without Shutdown) should not cause errors.
|
||
|
|
sub := &stubSubsystem{name: "plain-sub"}
|
||
|
|
svc, err := New(WithSubsystem(sub))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("New() failed: %v", err)
|
||
|
|
}
|
||
|
|
if err := svc.Shutdown(context.Background()); err != nil {
|
||
|
|
t.Fatalf("Shutdown() should succeed for non-shutdown subsystem: %v", err)
|
||
|
|
}
|
||
|
|
}
|