package tools import ( "context" "encoding/json" "github.com/mark3labs/mcp-go/mcp" "net/http" "net/http/httptest" "spark-mcp-go/internal/analyzer" "spark-mcp-go/internal/cluster" "spark-mcp-go/internal/httpclient" "spark-mcp-go/internal/storage" "strings" "testing" "time" ) func testDepsWithCluster(t *testing.T, rmURL, shsURL string) *Deps { 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: rmURL, SHSURL: shsURL, 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, DataDir: t.TempDir(), AnalyzerThresholds: analyzer.Thresholds{DataSkewRatio: 3.0, GCPressureRatio: 0.1, BottleneckShuffleGB: 50.0}} } func TestFetchSparkMetrics_Summary(t *testing.T) { shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/v1/applications/app_123/executors": w.Header().Set("Content-Type", "application/json") w.Write([]byte(`[{"id":"1","gc_time_ms":5000,"cpu_time_ms":10000}]`)) case "/api/v1/applications/app_123/stages": w.Header().Set("Content-Type", "application/json") w.Write([]byte(`[{"stageId":1,"duration_ms":30000,"shuffle_read_bytes":0,"shuffle_write_bytes":0,"max_partition_bytes":100,"min_partition_bytes":10}]`)) default: t.Errorf("unexpected SHS path %q", r.URL.Path) } })) defer shsSrv.Close() deps := testDepsWithCluster(t, "http://unused", shsSrv.URL) req := newToolRequest("fetch_spark_metrics", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123", "format": "summary"}) res, err := deps.FetchSparkMetricsHandler(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) } findings, ok := payload["findings"].([]any) if !ok || len(findings) == 0 { t.Fatalf("expected findings, got %+v", payload["findings"]) } } func TestFetchClusterEnv(t *testing.T) { rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/ws/v1/cluster/info": w.Header().Set("Content-Type", "application/json") w.Write([]byte(`{"clusterInfo":{"id":"rm1","name":"test"}}`)) case "/ws/v1/cluster/metrics": w.Header().Set("Content-Type", "application/json") w.Write([]byte(`{"clusterMetrics":{"appsSubmitted":5}}`)) default: t.Errorf("unexpected RM path %q", r.URL.Path) } })) defer rmSrv.Close() deps := testDepsWithCluster(t, rmSrv.URL, "http://unused") req := newToolRequest("fetch_cluster_env", map[string]any{"cluster_id": "cluster-a"}) res, err := deps.FetchClusterEnvHandler(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["cluster_info"] == nil { t.Errorf("cluster_info missing") } if payload["metrics"] == nil { t.Errorf("metrics missing") } if _, ok := payload["fetched_at"].(string); !ok { t.Errorf("fetched_at missing or not string") } } func TestAnalyzeSparkLog(t *testing.T) { rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/ws/v1/cluster/apps/app_123/amContainerLogs": w.WriteHeader(http.StatusNotFound) case "/ws/v1/cluster/apps/app_123/aggregated-logs": w.Write([]byte("first line\nsecond line\nERROR: something")) default: t.Errorf("unexpected RM path %q", r.URL.Path) } })) defer rmSrv.Close() shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/v1/applications/app_123/executors": w.Write([]byte(`[]`)) case "/api/v1/applications/app_123/stages": w.Write([]byte(`[]`)) default: t.Errorf("unexpected SHS path %q", r.URL.Path) } })) defer shsSrv.Close() deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL) req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"}) res, err := deps.AnalyzeSparkLogHandler(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) } prompt := payload["prompt"].(string) if !strings.Contains(prompt, "Spark Log Analysis") { t.Errorf("prompt missing header") } if !strings.Contains(payload["log_tail"].(string), "ERROR: something") { t.Errorf("log tail missing content") } } func TestAnalyzeSparkLog_LogSourceReported(t *testing.T) { rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/ws/v1/cluster/apps/app_123/amContainerLogs": w.WriteHeader(http.StatusNotFound) case "/ws/v1/cluster/apps/app_123/aggregated-logs": w.Write([]byte("log body")) default: t.Errorf("unexpected RM path %q", r.URL.Path) } })) defer rmSrv.Close() shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Write([]byte(`[]`)) })) defer shsSrv.Close() deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL) req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"}) res, err := deps.AnalyzeSparkLogHandler(context.Background(), req) if err != nil { t.Fatalf("handler error: %v", err) } if res.IsError { t.Fatalf("unexpected error result: %v", res.Content) } text, _ := mcp.AsTextContent(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["log_source"] == "" { t.Errorf("log_source empty") } } func TestAnalyzeSparkLog_PromptIncludesFindings(t *testing.T) { rmSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/ws/v1/cluster/apps/app_123/amContainerLogs": w.WriteHeader(http.StatusNotFound) case "/ws/v1/cluster/apps/app_123/aggregated-logs": w.Write([]byte("log body")) default: t.Errorf("unexpected RM path %q", r.URL.Path) } })) defer rmSrv.Close() shsSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/v1/applications/app_123/executors": w.Write([]byte(`[{"id":"1","gc_time_ms":5000,"cpu_time_ms":10000}]`)) case "/api/v1/applications/app_123/stages": w.Write([]byte(`[]`)) default: t.Errorf("unexpected SHS path %q", r.URL.Path) } })) defer shsSrv.Close() deps := testDepsWithCluster(t, rmSrv.URL, shsSrv.URL) req := newToolRequest("analyze_spark_log", map[string]any{"cluster_id": "cluster-a", "app_id": "app_123"}) res, err := deps.AnalyzeSparkLogHandler(context.Background(), req) if err != nil { t.Fatalf("handler error: %v", err) } if res.IsError { t.Fatalf("unexpected error result: %v", res.Content) } text, _ := mcp.AsTextContent(res.Content[0]) var payload map[string]any if err := json.Unmarshal([]byte(text.Text), &payload); err != nil { t.Fatalf("unmarshal result: %v", err) } prompt := payload["prompt"].(string) if !strings.Contains(prompt, "GC 压力") { t.Errorf("prompt missing GC finding; prompt:\n%s", prompt) } findings, ok := payload["findings"].([]any) if !ok || len(findings) == 0 { t.Fatalf("expected findings in result") } }