251 lines
8.4 KiB
Go
251 lines
8.4 KiB
Go
|
package mcpserver
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"email-mcp/internal/imapclient"
|
||
|
|
"email-mcp/internal/secretstore"
|
||
|
|
"email-mcp/internal/secretstore/kwallet"
|
||
|
|
)
|
||
|
|
|
||
|
|
type storeStub struct {
|
||
|
|
credential secretstore.Credential
|
||
|
|
loadErr error
|
||
|
|
loadCalls int
|
||
|
|
loadKey string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *storeStub) Save(context.Context, string, secretstore.Credential) error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *storeStub) Load(_ context.Context, key string) (secretstore.Credential, error) {
|
||
|
|
s.loadCalls++
|
||
|
|
s.loadKey = key
|
||
|
|
if s.loadErr != nil {
|
||
|
|
return secretstore.Credential{}, s.loadErr
|
||
|
|
}
|
||
|
|
return s.credential, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type serviceStub struct {
|
||
|
|
listMailboxes func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error)
|
||
|
|
listMessages func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error)
|
||
|
|
getMessage func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s serviceStub) ListMailboxes(ctx context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
return s.listMailboxes(ctx, cred)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s serviceStub) ListMessages(ctx context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
|
||
|
|
return s.listMessages(ctx, cred, mailbox, limit)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s serviceStub) GetMessage(ctx context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) {
|
||
|
|
return s.getMessage(ctx, cred, mailbox, uid)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerListMailboxesLoadsCredentialAndDelegates(t *testing.T) {
|
||
|
|
store := &storeStub{
|
||
|
|
credential: secretstore.Credential{
|
||
|
|
Host: "imap.example.com",
|
||
|
|
Username: "alice",
|
||
|
|
Password: "secret",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
server := New(store, serviceStub{
|
||
|
|
listMailboxes: func(_ context.Context, cred secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
if cred.Host != "imap.example.com" || cred.Username != "alice" || cred.Password != "secret" {
|
||
|
|
t.Fatalf("unexpected credential: %#v", cred)
|
||
|
|
}
|
||
|
|
return []imapclient.Mailbox{{Name: "INBOX"}}, nil
|
||
|
|
},
|
||
|
|
listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
|
||
|
|
t.Fatal("ListMessages should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
|
||
|
|
t.Fatal("GetMessage should not be called")
|
||
|
|
return imapclient.Message{}, nil
|
||
|
|
},
|
||
|
|
})
|
||
|
|
|
||
|
|
result, err := server.ListMailboxes(context.Background())
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("ListMailboxes returned error: %v", err)
|
||
|
|
}
|
||
|
|
if store.loadCalls != 1 {
|
||
|
|
t.Fatalf("expected credential to be loaded once, got %d", store.loadCalls)
|
||
|
|
}
|
||
|
|
if store.loadKey != secretstore.DefaultAccountKey {
|
||
|
|
t.Fatalf("expected load key %q, got %q", secretstore.DefaultAccountKey, store.loadKey)
|
||
|
|
}
|
||
|
|
if len(result) != 1 || result[0].Name != "INBOX" {
|
||
|
|
t.Fatalf("unexpected result: %#v", result)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerListMessagesLoadsCredentialAndDelegates(t *testing.T) {
|
||
|
|
store := &storeStub{
|
||
|
|
credential: secretstore.Credential{
|
||
|
|
Host: "imap.example.com",
|
||
|
|
Username: "alice",
|
||
|
|
Password: "secret",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
server := New(store, serviceStub{
|
||
|
|
listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
t.Fatal("ListMailboxes should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
|
||
|
|
if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 {
|
||
|
|
t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit)
|
||
|
|
}
|
||
|
|
return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil
|
||
|
|
},
|
||
|
|
getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
|
||
|
|
t.Fatal("GetMessage should not be called")
|
||
|
|
return imapclient.Message{}, nil
|
||
|
|
},
|
||
|
|
})
|
||
|
|
|
||
|
|
result, err := server.ListMessages(context.Background(), "INBOX", 5)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("ListMessages returned error: %v", err)
|
||
|
|
}
|
||
|
|
if len(result) != 1 || result[0].UID != 42 {
|
||
|
|
t.Fatalf("unexpected result: %#v", result)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestServerGetMessageUsesUIDContract(t *testing.T) {
|
||
|
|
store := &storeStub{
|
||
|
|
credential: secretstore.Credential{
|
||
|
|
Host: "imap.example.com",
|
||
|
|
Username: "alice",
|
||
|
|
Password: "secret",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
server := New(store, serviceStub{
|
||
|
|
listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
t.Fatal("ListMailboxes should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
|
||
|
|
t.Fatal("ListMessages should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
getMessage: func(_ context.Context, cred secretstore.Credential, mailbox string, uid uint32) (imapclient.Message, error) {
|
||
|
|
if cred.Host != "imap.example.com" || mailbox != "INBOX" || uid != 42 {
|
||
|
|
t.Fatalf("unexpected call: cred=%#v mailbox=%q uid=%d", cred, mailbox, uid)
|
||
|
|
}
|
||
|
|
return imapclient.Message{
|
||
|
|
UID: 42,
|
||
|
|
Mailbox: "INBOX",
|
||
|
|
Body: "body",
|
||
|
|
}, nil
|
||
|
|
},
|
||
|
|
})
|
||
|
|
|
||
|
|
message, err := server.GetMessage(context.Background(), "INBOX", 42)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatalf("GetMessage returned error: %v", err)
|
||
|
|
}
|
||
|
|
if message.UID != 42 || message.Mailbox != "INBOX" {
|
||
|
|
t.Fatalf("unexpected message: %#v", message)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRunnerRunWritesToolManifestAndHandlesRequests(t *testing.T) {
|
||
|
|
store := &storeStub{
|
||
|
|
credential: secretstore.Credential{
|
||
|
|
Host: "imap.example.com",
|
||
|
|
Username: "alice",
|
||
|
|
Password: "secret",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
input := bytes.NewBufferString("{\"tool\":\"list_messages\",\"arguments\":{\"mailbox\":\"INBOX\",\"limit\":5}}\n")
|
||
|
|
output := &bytes.Buffer{}
|
||
|
|
runner := NewRunner(New(store, serviceStub{
|
||
|
|
listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
t.Fatal("ListMailboxes should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
listMessages: func(_ context.Context, cred secretstore.Credential, mailbox string, limit int) ([]imapclient.MessageSummary, error) {
|
||
|
|
if cred.Host != "imap.example.com" || mailbox != "INBOX" || limit != 5 {
|
||
|
|
t.Fatalf("unexpected call: cred=%#v mailbox=%q limit=%d", cred, mailbox, limit)
|
||
|
|
}
|
||
|
|
return []imapclient.MessageSummary{{UID: 42, Subject: "hello", From: "alice@example.com"}}, nil
|
||
|
|
},
|
||
|
|
getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
|
||
|
|
t.Fatal("GetMessage should not be called")
|
||
|
|
return imapclient.Message{}, nil
|
||
|
|
},
|
||
|
|
}), input, output, &bytes.Buffer{})
|
||
|
|
|
||
|
|
if err := runner.Run(context.Background()); err != nil {
|
||
|
|
t.Fatalf("Run returned error: %v", err)
|
||
|
|
}
|
||
|
|
if store.loadCalls != 1 {
|
||
|
|
t.Fatalf("expected credential preload once, got %d", store.loadCalls)
|
||
|
|
}
|
||
|
|
|
||
|
|
decoder := json.NewDecoder(output)
|
||
|
|
|
||
|
|
var manifest struct {
|
||
|
|
Tools []struct {
|
||
|
|
Name string `json:"name"`
|
||
|
|
} `json:"tools"`
|
||
|
|
}
|
||
|
|
if err := decoder.Decode(&manifest); err != nil {
|
||
|
|
t.Fatalf("failed to decode manifest: %v", err)
|
||
|
|
}
|
||
|
|
if len(manifest.Tools) != 3 {
|
||
|
|
t.Fatalf("expected 3 tools, got %#v", manifest.Tools)
|
||
|
|
}
|
||
|
|
if manifest.Tools[0].Name != "list_mailboxes" || manifest.Tools[1].Name != "list_messages" || manifest.Tools[2].Name != "get_message" {
|
||
|
|
t.Fatalf("unexpected tool manifest: %#v", manifest.Tools)
|
||
|
|
}
|
||
|
|
|
||
|
|
var response struct {
|
||
|
|
Result []imapclient.MessageSummary `json:"result"`
|
||
|
|
}
|
||
|
|
if err := decoder.Decode(&response); err != nil {
|
||
|
|
t.Fatalf("failed to decode response: %v", err)
|
||
|
|
}
|
||
|
|
if len(response.Result) != 1 || response.Result[0].UID != 42 {
|
||
|
|
t.Fatalf("unexpected response: %#v", response.Result)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestRunnerRunReturnsFriendlyMissingCredentialError(t *testing.T) {
|
||
|
|
store := &storeStub{
|
||
|
|
loadErr: kwallet.ErrCredentialNotFound,
|
||
|
|
}
|
||
|
|
runner := NewRunner(New(store, serviceStub{
|
||
|
|
listMailboxes: func(context.Context, secretstore.Credential) ([]imapclient.Mailbox, error) {
|
||
|
|
t.Fatal("ListMailboxes should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
listMessages: func(context.Context, secretstore.Credential, string, int) ([]imapclient.MessageSummary, error) {
|
||
|
|
t.Fatal("ListMessages should not be called")
|
||
|
|
return nil, nil
|
||
|
|
},
|
||
|
|
getMessage: func(context.Context, secretstore.Credential, string, uint32) (imapclient.Message, error) {
|
||
|
|
t.Fatal("GetMessage should not be called")
|
||
|
|
return imapclient.Message{}, nil
|
||
|
|
},
|
||
|
|
}), bytes.NewBuffer(nil), &bytes.Buffer{}, &bytes.Buffer{})
|
||
|
|
|
||
|
|
err := runner.Run(context.Background())
|
||
|
|
if !errors.Is(err, ErrCredentialsNotConfigured) {
|
||
|
|
t.Fatalf("expected missing credential error, got %v", err)
|
||
|
|
}
|
||
|
|
}
|