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

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