wenshu/k8s-orchestrator/test/api/server_test.go

package api_test

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

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

func TestAddNodeRequiresIdempotencyKey(t *testing.T) {
	h := newTestServer(t)
	body := []byte(`{"cpu":8,"ram":32,"ssd":512,"ipAddress":"10.0.0.2","sshPassword":"x"}`)
	req := httptest.NewRequest(http.MethodPost, "/v1/nodes", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp := httptest.NewRecorder()

	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", resp.Code)
	}
}

func TestAddNodeIdempotent(t *testing.T) {
	h := newTestServer(t)
	payload := map[string]interface{}{
		"cpu":         8,
		"ram":         32,
		"ssd":         512,
		"ipAddress":   "10.0.0.2",
		"sshPassword": "x",
	}
	body, _ := json.Marshal(payload)

	first := httptest.NewRequest(http.MethodPost, "/v1/nodes", bytes.NewReader(body))
	first.Header.Set("Content-Type", "application/json")
	first.Header.Set("Idempotency-Key", "node-op-1")
	resp1 := httptest.NewRecorder()
	h.ServeHTTP(resp1, first)
	if resp1.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp1.Code)
	}

	second := httptest.NewRequest(http.MethodPost, "/v1/nodes", bytes.NewReader(body))
	second.Header.Set("Content-Type", "application/json")
	second.Header.Set("Idempotency-Key", "node-op-1")
	resp2 := httptest.NewRecorder()
	h.ServeHTTP(resp2, second)
	if resp2.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp2.Code)
	}

	var p1 map[string]string
	var p2 map[string]string
	_ = json.Unmarshal(resp1.Body.Bytes(), &p1)
	_ = json.Unmarshal(resp2.Body.Bytes(), &p2)
	if p1["operationId"] == "" || p2["operationId"] == "" {
		t.Fatalf("operationId should not be empty")
	}
	if p1["operationId"] != p2["operationId"] {
		t.Fatalf("expected idempotent same operationId, got %s and %s", p1["operationId"], p2["operationId"])
	}
}

func TestCreateStarRocksClusterRequiresIdempotencyKey(t *testing.T) {
	h := newTestServer(t)
	body := []byte(`{"feNodeIds":["node-a"],"beNodeIds":["node-b"]}`)
	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp := httptest.NewRecorder()

	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", resp.Code)
	}
}

func TestCreateStarRocksClusterIdempotent(t *testing.T) {
	h := newTestServer(t)
	body := []byte(`{"clusterId":"sr-1","feNodeIds":["node-a"],"beNodeIds":["node-b","node-c"]}`)

	first := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters", bytes.NewReader(body))
	first.Header.Set("Content-Type", "application/json")
	first.Header.Set("Idempotency-Key", "sr-create-1")
	resp1 := httptest.NewRecorder()
	h.ServeHTTP(resp1, first)
	if resp1.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp1.Code)
	}

	second := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters", bytes.NewReader(body))
	second.Header.Set("Content-Type", "application/json")
	second.Header.Set("Idempotency-Key", "sr-create-1")
	resp2 := httptest.NewRecorder()
	h.ServeHTTP(resp2, second)
	if resp2.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp2.Code)
	}

	var p1 map[string]string
	var p2 map[string]string
	_ = json.Unmarshal(resp1.Body.Bytes(), &p1)
	_ = json.Unmarshal(resp2.Body.Bytes(), &p2)
	if p1["operationId"] == "" || p2["operationId"] == "" {
		t.Fatalf("operationId should not be empty")
	}
	if p1["operationId"] != p2["operationId"] {
		t.Fatalf("expected idempotent same operationId, got %s and %s", p1["operationId"], p2["operationId"])
	}
}

func TestGetStarRocksCluster(t *testing.T) {
	cluster := &orcv1.StarRocksVirtualCluster{
		ObjectMeta: metav1.ObjectMeta{Name: "sr-get"},
		Spec: orcv1.StarRocksVirtualClusterSpec{
			FeNodeRefs: []string{"node-a"},
			BeNodeRefs: []string{"node-b"},
		},
		Status: orcv1.StarRocksVirtualClusterStatus{
			Phase:     orcv1.StarRocksClusterPhaseReady,
			Namespace: "starrocks-sr-get",
		},
	}
	h := newTestServerWithObjects(t, cluster)

	req := httptest.NewRequest(http.MethodGet, "/v1/starrocks-clusters/sr-get", nil)
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", resp.Code)
	}

	var got orcv1.StarRocksVirtualCluster
	if err := json.Unmarshal(resp.Body.Bytes(), &got); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if got.Name != "sr-get" {
		t.Fatalf("expected cluster name sr-get, got %s", got.Name)
	}
}

func TestReleaseStarRocksClusterAccepted(t *testing.T) {
	h := newTestServer(t)

	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters/sr-1/release", nil)
	req.Header.Set("Idempotency-Key", "sr-release-1")
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp.Code)
	}
}

func TestScaleOutStarRocksClusterRequiresNodeIDs(t *testing.T) {
	h := newTestServer(t)

	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters/sr-1/scale-out", bytes.NewReader([]byte(`{}`)))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Idempotency-Key", "sr-scale-out-1")
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", resp.Code)
	}
}

func TestScaleInStarRocksClusterAccepted(t *testing.T) {
	h := newTestServer(t)

	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-clusters/sr-1/scale-in", bytes.NewReader([]byte(`{"nodeIds":["node-b"]}`)))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Idempotency-Key", "sr-scale-in-1")
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp.Code)
	}
}

func TestSubmitStarRocksJobRequiresIdempotencyKey(t *testing.T) {
	h := newTestServer(t)
	body := []byte(`{"clusterId":"sr-1","jobSpec":{"template":{"spec":{"restartPolicy":"Never","containers":[{"name":"task","image":"busybox","command":["sh","-c","echo ok"]}]}}}}`)
	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-jobs", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp := httptest.NewRecorder()

	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusBadRequest {
		t.Fatalf("expected 400, got %d", resp.Code)
	}
}

func TestSubmitStarRocksJobAccepted(t *testing.T) {
	h := newTestServer(t)
	body := []byte(`{"clusterId":"sr-1","jobSpec":{"template":{"spec":{"restartPolicy":"Never","containers":[{"name":"task","image":"busybox","command":["sh","-c","echo ok"]}]}}}}`)
	req := httptest.NewRequest(http.MethodPost, "/v1/starrocks-jobs", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Idempotency-Key", "sr-job-1")
	resp := httptest.NewRecorder()

	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusAccepted {
		t.Fatalf("expected 202, got %d", resp.Code)
	}
}

func TestGetStarRocksJobStatus(t *testing.T) {
	op := &orcv1.ClusterOperation{
		ObjectMeta: metav1.ObjectMeta{Name: "op-submit-starrocks-1"},
		Spec: orcv1.ClusterOperationSpec{
			Type: orcv1.OperationSubmitStarRocksJob,
			SubmitStarRocks: &orcv1.SubmitStarRocksJobSpec{
				ClusterName: "sr-1",
			},
		},
		Status: orcv1.ClusterOperationStatus{
			Phase:   orcv1.OperationPhaseRunning,
			Message: "submitted",
		},
	}
	h := newTestServerWithObjects(t, op)

	req := httptest.NewRequest(http.MethodGet, "/v1/starrocks-jobs/operations/op-submit-starrocks-1", nil)
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d, body=%s", resp.Code, resp.Body.String())
	}

	var payload map[string]interface{}
	if err := json.Unmarshal(resp.Body.Bytes(), &payload); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if payload["operationId"] != "op-submit-starrocks-1" {
		t.Fatalf("unexpected operationId: %v", payload["operationId"])
	}
}

func newTestServer(t *testing.T) http.Handler {
	t.Helper()
	return newTestServerWithObjects(t)
}