spark-mcp/internal/storage/submission_repo_test.go

package storage

import (
	"context"
	"testing"
	"time"
)

func setupSubmissionTest(t *testing.T) (*SubmissionRepo, *UploadRepo, context.Context) {
	db, err := Open(":memory:")
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	return db.Submissions(), db.Uploads(), context.Background()
}

func TestSubmissionRepo_RoundTrip(t *testing.T) {
	repo, uploadRepo, ctx := setupSubmissionTest(t)

	if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "alpha.py", 10, "deadbeef", time.Now()); err != nil {
		t.Fatalf("create upload: %v", err)
	}

	s := Submission{
		FileID:      "00000000000000000000000000000001",
		AppID:       "application_1_0001",
		TrackingURL: "http://rm.example.com:8088/proxy/application_1_0001/",
		ClusterID:   "cluster-a",
		ExitCode:    0,
		DurationMS:  1234,
		SubmittedAt: time.Unix(0, time.Now().UnixNano()),
	}
	if err := repo.Create(ctx, s); err != nil {
		t.Fatalf("create: %v", err)
	}

	list, err := repo.List(ctx, SubmissionFilter{Limit: 10})
	if err != nil {
		t.Fatalf("list: %v", err)
	}
	if len(list) != 1 {
		t.Fatalf("len(list)=%d, want 1", len(list))
	}
	swf := list[0]
	if swf.FileName != "alpha.py" {
		t.Errorf("file_name=%q, want alpha.py", swf.FileName)
	}
	if swf.AppID != s.AppID {
		t.Errorf("app_id=%q, want %q", swf.AppID, s.AppID)
	}

	id := swf.ID
	got, err := repo.Get(ctx, id)
	if err != nil {
		t.Fatalf("get: %v", err)
	}
	if got.ID != id || got.AppID != s.AppID {
		t.Errorf("got=%+v, want id=%d app_id=%s", got, id, s.AppID)
	}

	if err := repo.Delete(ctx, id); err != nil {
		t.Fatalf("delete: %v", err)
	}
	if _, err := repo.Get(ctx, id); !isErrNotFound(err) {
		t.Errorf("expected ErrNotFound after delete, got %v", err)
	}
}

func TestSubmissionRepo_SearchByAppID(t *testing.T) {
	repo, uploadRepo, ctx := setupSubmissionTest(t)

	if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "alpha.py", 10, "deadbeef", time.Now()); err != nil {
		t.Fatalf("create upload alpha: %v", err)
	}
	if err := uploadRepo.Create(ctx, "00000000000000000000000000000002", "beta.py", 10, "cafebabe", time.Now()); err != nil {
		t.Fatalf("create upload beta: %v", err)
	}
	if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000001", AppID: "application_alpha_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
		t.Fatalf("create alpha: %v", err)
	}
	if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000002", AppID: "application_beta_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
		t.Fatalf("create beta: %v", err)
	}

	list, err := repo.List(ctx, SubmissionFilter{Search: "alpha", Limit: 10})
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(list) != 1 || list[0].AppID != "application_alpha_1" || list[0].FileName != "alpha.py" {
		t.Errorf("got %+v, want 1 alpha submission", list)
	}
}

func TestSubmissionRepo_SearchByFileName(t *testing.T) {
	db, err := Open(":memory:")
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	defer db.Close()
	repo := db.Submissions()
	uploadRepo := db.Uploads()
	ctx := context.Background()

	if err := uploadRepo.Create(ctx, "00000000000000000000000000000001", "report.py", 10, "deadbeef", time.Now()); err != nil {
		t.Fatalf("create upload: %v", err)
	}
	if err := repo.Create(ctx, Submission{FileID: "00000000000000000000000000000001", AppID: "application_x_1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
		t.Fatalf("create: %v", err)
	}

	list, err := repo.List(ctx, SubmissionFilter{Search: "report", Limit: 10})
	if err != nil {
		t.Fatalf("search: %v", err)
	}
	if len(list) != 1 || list[0].FileName != "report.py" {
		t.Errorf("got %+v, want 1 report submission", list)
	}
}

func TestSubmissionRepo_FilterByFileID(t *testing.T) {
	db, err := Open(":memory:")
	if err != nil {
		t.Fatalf("open: %v", err)
	}
	defer db.Close()
	repo := db.Submissions()
	ctx := context.Background()

	if err := repo.Create(ctx, Submission{FileID: "f1", AppID: "a1", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
		t.Fatalf("create f1: %v", err)
	}
	if err := repo.Create(ctx, Submission{FileID: "f2", AppID: "a2", ClusterID: "c", ExitCode: 0, DurationMS: 1, SubmittedAt: time.Now()}); err != nil {
		t.Fatalf("create f2: %v", err)
	}

	list, err := repo.List(ctx, SubmissionFilter{FileID: "f1", Limit: 10})
	if err != nil {
		t.Fatalf("list: %v", err)
	}
	if len(list) != 1 || list[0].FileID != "f1" {
		t.Errorf("got %+v, want f1 submission", list)
	}
}

func isErrNotFound(err error) bool {
	if err == nil {
		return false
	}
	return err.Error() == ErrNotFound.Error()
}