wenshu/k8s-orchestrator/internal/controller/sparkapplication_cleaner_test.go

package controller

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
	ctrlclientfake "sigs.k8s.io/controller-runtime/pkg/client/fake"

	orcv1 "github.com/cyrus/spark-orchestrator/api/v1alpha1"
	"github.com/cyrus/spark-orchestrator/internal/k8s"
	"github.com/cyrus/spark-orchestrator/internal/util"
)

func TestSparkApplicationCleanerReconciler_cleanupNamespacePushesBeforeDelete(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(orcv1.AddToScheme(scheme))

	terminationTime := time.Now().Add(-2 * time.Hour).UTC()
	app := newTerminalSparkApplication("spark-ns", "spark-app", terminationTime)
	op := newSparkSubmitOperation("op-1", "spark-ns/spark-app")

	var capturedBody []byte
	var capturedHeader string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			t.Fatalf("unexpected method: %s", r.Method)
		}
		if r.URL.Path != "/api/data/internal/task-exec/spark-jobs/op-1/archive" {
			t.Fatalf("unexpected path: %s", r.URL.Path)
		}
		capturedHeader = r.Header.Get("X-Internal-Token")
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("read body: %v", err)
		}
		capturedBody = append([]byte(nil), body...)
		w.WriteHeader(http.StatusOK)
	}))
	defer server.Close()

	client := ctrlclientfake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(op, app).
		Build()

	reconciler := &SparkApplicationCleanerReconciler{
		Client:                      client,
		TTLSeconds:                  int64(time.Hour.Seconds()),
		ArchiveGraceSeconds:         int64(time.Hour.Seconds()),
		BackendBaseURL:              server.URL,
		BackendTimeout:              5 * time.Second,
		BackendInternalToken:        "internal-secret",
		BackendInternalHeader:       "X-Internal-Token",
		OrchestratorSecretNamespace: "orchestrator-system",
	}

	deleted, err := reconciler.cleanupNamespace(context.Background(), "spark-ns")
	if err != nil {
		t.Fatalf("cleanupNamespace: %v", err)
	}
	if deleted != 1 {
		t.Fatalf("expected one deletion, got %d", deleted)
	}
	if capturedHeader != "internal-secret" {
		t.Fatalf("expected internal token header, got %q", capturedHeader)
	}
	if len(capturedBody) == 0 {
		t.Fatalf("expected archive body")
	}
	var payload map[string]any
	if err := json.Unmarshal(capturedBody, &payload); err != nil {
		t.Fatalf("unmarshal body: %v", err)
	}
	if payload["operationId"] != "op-1" {
		t.Fatalf("unexpected operationId: %v", payload["operationId"])
	}

	var remaining unstructured.Unstructured
	remaining.SetAPIVersion(k8s.SparkApplicationAPIVersion)
	remaining.SetKind(k8s.SparkApplicationKind)
	err = client.Get(context.Background(), ctrlclient.ObjectKey{Name: "spark-app", Namespace: "spark-ns"}, &remaining)
	if err == nil {
		t.Fatalf("expected spark application to be deleted")
	}
}

func TestSparkApplicationCleanerReconciler_cleanupNamespaceRetainsAppOnPushFailureUntilGraceExpires(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(orcv1.AddToScheme(scheme))

	terminationTime := time.Now().Add(-70 * time.Minute).UTC()
	app := newTerminalSparkApplication("spark-ns", "spark-app", terminationTime)
	op := newSparkSubmitOperation("op-2", "spark-ns/spark-app")

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusServiceUnavailable)
	}))
	defer server.Close()

	client := ctrlclientfake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(op, app).
		Build()

	reconciler := &SparkApplicationCleanerReconciler{
		Client:                      client,
		TTLSeconds:                  int64(time.Hour.Seconds()),
		ArchiveGraceSeconds:         int64(time.Hour.Seconds()),
		BackendBaseURL:              server.URL,
		BackendTimeout:              5 * time.Second,
		BackendInternalHeader:       "X-Internal-Token",
		BackendInternalToken:        "internal-secret",
		OrchestratorSecretNamespace: "orchestrator-system",
	}

	deleted, err := reconciler.cleanupNamespace(context.Background(), "spark-ns")
	if err != nil {
		t.Fatalf("cleanupNamespace: %v", err)
	}
	if deleted != 0 {
		t.Fatalf("expected no deletion while grace is active, got %d", deleted)
	}

	var remaining unstructured.Unstructured
	remaining.SetAPIVersion(k8s.SparkApplicationAPIVersion)
	remaining.SetKind(k8s.SparkApplicationKind)
	if err := client.Get(context.Background(), ctrlclient.ObjectKey{Name: "spark-app", Namespace: "spark-ns"}, &remaining); err != nil {
		t.Fatalf("expected spark application to remain: %v", err)
	}
}

func TestSparkApplicationCleanerReconciler_cleanupNamespaceForceDeletesAfterGrace(t *testing.T) {
	scheme := runtime.NewScheme()
	utilruntime.Must(orcv1.AddToScheme(scheme))

	terminationTime := time.Now().Add(-2*time.Hour - 5*time.Minute).UTC()
	app := newTerminalSparkApplication("spark-ns", "spark-app", terminationTime)
	op := newSparkSubmitOperation("op-3", "spark-ns/spark-app")

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusBadGateway)
	}))
	defer server.Close()

	client := ctrlclientfake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(op, app).
		Build()

	reconciler := &SparkApplicationCleanerReconciler{
		Client:                      client,
		TTLSeconds:                  int64(time.Hour.Seconds()),
		ArchiveGraceSeconds:         int64(time.Hour.Seconds()),
		BackendBaseURL:              server.URL,
		BackendTimeout:              5 * time.Second,
		BackendInternalHeader:       "X-Internal-Token",
		BackendInternalToken:        "internal-secret",
		OrchestratorSecretNamespace: "orchestrator-system",
	}

	deleted, err := reconciler.cleanupNamespace(context.Background(), "spark-ns")
	if err != nil {
		t.Fatalf("cleanupNamespace: %v", err)
	}
	if deleted != 1 {
		t.Fatalf("expected forced deletion after grace, got %d", deleted)
	}

	var remaining unstructured.Unstructured
	remaining.SetAPIVersion(k8s.SparkApplicationAPIVersion)
	remaining.SetKind(k8s.SparkApplicationKind)
	if err := client.Get(context.Background(), ctrlclient.ObjectKey{Name: "spark-app", Namespace: "spark-ns"}, &remaining); err == nil {
		t.Fatalf("expected spark application to be deleted")
	}
}

func newSparkSubmitOperation(name, resourceRef string) *orcv1.ClusterOperation {
	return &orcv1.ClusterOperation{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: map[string]string{util.LabelOperationType: string(orcv1.OperationSubmitSparkJob)},
		},
		Spec: orcv1.ClusterOperationSpec{
			Type: orcv1.OperationSubmitSparkJob,
		},
		Status: orcv1.ClusterOperationStatus{
			ResourceRef: resourceRef,
		},
	}
}

func newTerminalSparkApplication(namespace, name string, terminationTime time.Time) *unstructured.Unstructured {
	return &unstructured.Unstructured{
		Object: map[string]interface{}{
			"apiVersion": k8s.SparkApplicationAPIVersion,
			"kind":       k8s.SparkApplicationKind,
			"metadata": map[string]interface{}{
				"name":      name,
				"namespace": namespace,
				"labels": map[string]interface{}{
					util.LabelManagedBy: util.ManagedByValue,
				},
			},
			"status": map[string]interface{}{
				"terminationTime": terminationTime.Format(time.RFC3339),
				"applicationState": map[string]interface{}{
					"state":               k8s.SparkAppStateCompleted,
					"stateTransitionTime": terminationTime.Format(time.RFC3339),
				},
			},
		},
	}
}