spark-mcp/internal/mcp/tools/high_level_test.go

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")
	}
}