gopeed/pkg/mcpserver/speed_test.go

package mcpserver

import (
	"context"
	"net/http/httptest"
	"testing"

	"github.com/GopeedLab/gopeed/pkg/download"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func TestGetDownloadSpeedTool(t *testing.T) {
	downloader := download.NewDownloader(&download.DownloaderConfig{
		Storage:    download.NewMemStorage(),
		StorageDir: t.TempDir(),
	})
	if err := downloader.Setup(); err != nil {
		t.Fatal(err)
	}
	defer downloader.Close()

	httpServer := httptest.NewServer(NewHandler(downloader))
	defer httpServer.Close()

	client := mcp.NewClient(&mcp.Implementation{Name: "gopeed-test", Version: "1.0.0"}, nil)
	session, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil)
	if err != nil {
		t.Fatal(err)
	}
	defer session.Close()

	// 1. Query all tasks speed (empty list)
	result, err := session.CallTool(context.Background(), &mcp.CallToolParams{
		Name:      "get_download_speed",
		Arguments: map[string]any{},
	})
	if err != nil {
		t.Fatal(err)
	}
	if result.IsError {
		t.Fatalf("get_download_speed returned error: %+v", result.Content)
	}
	structured, ok := result.StructuredContent.(map[string]any)
	if !ok {
		t.Fatalf("structured content type = %T", result.StructuredContent)
	}
	if activeCount, ok := structured["activeTaskCount"].(float64); !ok || activeCount != 0 {
		t.Fatalf("activeTaskCount = %v, want 0", structured["activeTaskCount"])
	}

	// 2. Query non-existent task ID -> should return error
	notFoundResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
		Name: "get_download_speed",
		Arguments: map[string]any{
			"id": "non-existent-task-id",
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if !notFoundResult.IsError {
		t.Fatal("get_download_speed should return error for non-existent task ID")
	}
}