Phase 5 Batch 2: 4 RM Tool + rm 客户端 + 日志降级链

- internal/rm/client.go: YARN RM 客户端 - ListApps / GetApp / KillApp / GetLogs - GetLogs 降级链数据驱动 (amContainerLogs → aggregated-logs → logs) - 跨主机 redirect 保留 Authorization (URLAllowlist 兜底) - 4 Tool: list_applications / get_application_status / get_application_logs / kill_application - mcp.WithEnum(state) 约束 YARN app 状态 - tool handler 永远 (result, nil), 业务错误用 NewToolResultError - deps.go: +HTTPClient + MaxResponseBytes - main.go: 注入 httpclient.Client - 端到端实测: mock RM 验证 4 Tool + 降级链 + 错误路径 Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-10 16:50:27 +08:00
parent 705665745e
commit c4ac3cc354
9 changed files with 885 additions and 0 deletions
@@ -0,0 +1,176 @@
+package tools
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/mark3labs/mcp-go/mcp"
+
+	"spark-mcp-go/internal/cluster"
+	"spark-mcp-go/internal/httpclient"
+	"spark-mcp-go/internal/storage"
+)
+
+func testDeps(t *testing.T, srv *httptest.Server) (*Deps, *storage.ClusterRepo) {
+	t.Helper()
+	db, err := storage.Open(":memory:")
+	if err != nil {
+		t.Fatalf("open db: %v", err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+
+	cl := &cluster.Cluster{
+		ID:                    "cluster-a",
+		Name:                  "Cluster A",
+		RMURL:                 srv.URL,
+		SHSURL:                "http://shs.example.com:18080",
+		SparkSubmitExecuteBin: "/usr/bin/spark-submit",
+		IsActive:              true,
+		AuthType:              cluster.AuthNone,
+	}
+	if err := db.Clusters().Create(context.Background(), cl); err != nil {
+		t.Fatalf("create cluster: %v", err)
+	}
+
+	return &Deps{
+		HTTPClient: httpclient.New(httpclient.Config{
+			Timeout:          5 * time.Second,
+			MaxResponseBytes: 1 << 20,
+		}),
+		ClusterRepo:      db.Clusters(),
+		MaxResponseBytes: 1 << 20,
+	}, db.Clusters()
+}
+
+func newToolRequest(name string, args map[string]any) mcp.CallToolRequest {
+	return mcp.CallToolRequest{
+		Request: mcp.Request{Method: "tools/call"},
+		Params: mcp.CallToolParams{
+			Name:      name,
+			Arguments: args,
+		},
+	}
+}
+
+func TestListApplications_EndToEnd(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/ws/v1/cluster/apps" {
+			t.Errorf("path=%q", r.URL.Path)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		w.Write([]byte(`{"apps":{"app":[{"id":"application_1"}]}}`))
+	}))
+	defer srv.Close()
+
+	deps, _ := testDeps(t, srv)
+	req := newToolRequest("list_applications", map[string]any{
+		"cluster_id": "cluster-a",
+		"state":      "RUNNING",
+	})
+
+	res, err := deps.ListApplicationsHandler(context.Background(), req)
+	if err != nil {
+		t.Fatalf("handler error: %v", err)
+	}
+	if res.IsError {
+		t.Fatalf("unexpected error result: %v", res.Content)
+	}
+	text, ok := mcp.AsTextContent(res.Content[0])
+	if !ok {
+		t.Fatalf("content is not text: %T", res.Content[0])
+	}
+	if !strings.Contains(text.Text, "application_1") {
+		t.Errorf("result missing app: %s", text.Text)
+	}
+}
+
+func TestKillApplication_AppIDNotProvided(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Errorf("unexpected RM call")
+	}))
+	defer srv.Close()
+
+	deps, _ := testDeps(t, srv)
+	req := newToolRequest("kill_application", map[string]any{
+		"cluster_id": "cluster-a",
+	})
+
+	res, err := deps.KillApplicationHandler(context.Background(), req)
+	if err != nil {
+		t.Fatalf("handler error: %v", err)
+	}
+	if !res.IsError {
+		t.Errorf("expected error result, got: %v", res.Content)
+	}
+}
+
+func TestGetApplicationLogs_SourceReported(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/ws/v1/cluster/apps/app_123/amContainerLogs" {
+			w.Write([]byte("driver stdout"))
+			return
+		}
+		t.Errorf("unexpected path %q", r.URL.Path)
+	}))
+	defer srv.Close()
+
+	deps, _ := testDeps(t, srv)
+	req := newToolRequest("get_application_logs", map[string]any{
+		"cluster_id": "cluster-a",
+		"app_id":     "app_123",
+	})
+
+	res, err := deps.GetApplicationLogsHandler(context.Background(), req)
+	if err != nil {
+		t.Fatalf("handler error: %v", err)
+	}
+	if res.IsError {
+		t.Fatalf("unexpected error result: %v", res.Content)
+	}
+
+	text, ok := mcp.AsTextContent(res.Content[0])
+	if !ok {
+		t.Fatalf("content is not text: %T", res.Content[0])
+	}
+	var payload map[string]any
+	if err := json.Unmarshal([]byte(text.Text), &payload); err != nil {
+		t.Fatalf("unmarshal result: %v", err)
+	}
+	if payload["source"] != "amContainerLogs" {
+		t.Errorf("source=%v, want amContainerLogs", payload["source"])
+	}
+	if !strings.Contains(payload["text"].(string), "driver stdout") {
+		t.Errorf("text missing logs: %v", payload["text"])
+	}
+}
+
+func TestKillApplication_NotFound(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/ws/v1/cluster/apps/app_123/state" {
+			w.WriteHeader(http.StatusNotFound)
+			w.Write([]byte(`{"error":"Not Found"}`))
+			return
+		}
+		t.Errorf("unexpected path %q", r.URL.Path)
+	}))
+	defer srv.Close()
+
+	deps, _ := testDeps(t, srv)
+	req := newToolRequest("kill_application", map[string]any{
+		"cluster_id": "cluster-a",
+		"app_id":     "app_123",
+	})
+
+	res, err := deps.KillApplicationHandler(context.Background(), req)
+	if err != nil {
+		t.Fatalf("handler error: %v", err)
+	}
+	if !res.IsError {
+		t.Errorf("expected error result, got: %v", res.Content)
+	}
+}