easyai-ai-gateway/apps/api/internal/capacitycontroller/kubernetes_test.go

package capacitycontroller

import (
	"context"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestKubernetesClientReadsSiteAndMutatesOnlyWorkerScale(t *testing.T) {
	var patches []string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		if request.Header.Get("Authorization") != "Bearer test-token" {
			http.Error(w, "unauthorized", http.StatusUnauthorized)
			return
		}
		switch {
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-hongkong"):
			_, _ = w.Write([]byte(`{
				"spec":{"replicas":1,"template":{"spec":{"containers":[{
					"name":"worker","image":"registry.invalid/gateway@sha256:abc",
					"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
					"resources":{"requests":{"memory":"512Mi","cpu":"250m"}}
				}]}}}
			}`))
		case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
			if request.URL.Query().Get("labelSelector") != "easyai.io/site=hongkong,easyai.io/worker=true" {
				http.Error(w, "unexpected Worker node selector", http.StatusBadRequest)
				return
			}
			_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"easyai-hongkong"},"status":{
				"allocatable":{"memory":"8Gi","cpu":"4"},
				"conditions":[{"type":"MemoryPressure","status":"False"}]
			}}]}`))
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/api/v1/namespaces/easyai/pods"):
			_, _ = w.Write([]byte(`{"items":[{"metadata":{"name":"worker-hongkong-1"},"spec":{"nodeName":"easyai-hongkong"}}]}`))
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
			strings.HasSuffix(request.URL.Path, "/nodes/easyai-hongkong"):
			_, _ = w.Write([]byte(`{"usage":{"memory":"3Gi","cpu":"500m"}}`))
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/metrics.k8s.io/") &&
			strings.HasSuffix(request.URL.Path, "/pods"):
			_, _ = w.Write([]byte(`{"items":[{
				"metadata":{"name":"worker-hongkong-1"},
				"containers":[{"usage":{"memory":"384Mi","cpu":"125m"}}]
			}]}`))
		case request.Method == http.MethodPatch:
			patches = append(patches, request.URL.Path)
			w.Header().Set("Content-Type", "application/json")
			_, _ = w.Write([]byte(`{}`))
		default:
			http.NotFound(w, request)
		}
	}))
	defer server.Close()
	tokenFile := filepath.Join(t.TempDir(), "token")
	if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
		t.Fatal(err)
	}
	client, err := NewKubernetesClient(KubernetesConfig{
		Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	state, err := client.SiteState(context.Background(), "hongkong")
	if err != nil {
		t.Fatal(err)
	}
	if state.CurrentReplicas != 1 || state.NodeName != "easyai-hongkong" ||
		state.WorkerRequestMemoryBytes != 512<<20 || state.WorkerRequestMilliCPU != 250 ||
		state.AllocatableMemoryBytes != 8<<30 || state.UsedMemoryBytes != 3<<30 ||
		state.Nodes[0].WorkerUsedMemoryBytes != 384<<20 || state.Nodes[0].WorkerUsedMilliCPU != 125 {
		t.Fatalf("site state=%+v", state)
	}
	if err := client.SetPodDeletionCost(context.Background(), "worker-pod", -1000); err != nil {
		t.Fatal(err)
	}
	if err := client.ScaleWorkerDeployment(context.Background(), "hongkong", 2); err != nil {
		t.Fatal(err)
	}
	if len(patches) != 2 ||
		!strings.Contains(patches[0], "/pods/worker-pod") ||
		!strings.Contains(patches[1], "/deployments/easyai-worker-hongkong/scale") {
		t.Fatalf("patches=%v", patches)
	}
}

func TestKubernetesClientAllowsDisabledSiteWithoutWorkerNodes(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		switch {
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
			_, _ = w.Write([]byte(`{
				"spec":{"replicas":0,"template":{"spec":{"containers":[{
					"name":"worker","image":"registry.invalid/gateway@sha256:abc",
					"env":[{"name":"AI_GATEWAY_REVISION","value":"release-sha"}],
					"resources":{"requests":{"memory":"1536Mi","cpu":"500m"}}
				}]}}}
			}`))
		case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
			_, _ = w.Write([]byte(`{"items":[]}`))
		default:
			http.NotFound(w, request)
		}
	}))
	defer server.Close()
	tokenFile := filepath.Join(t.TempDir(), "token")
	if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
		t.Fatal(err)
	}
	client, err := NewKubernetesClient(KubernetesConfig{
		Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	state, err := client.SiteState(context.Background(), "ningbo")
	if err != nil {
		t.Fatal(err)
	}
	if state.Site != "ningbo" || state.CurrentReplicas != 0 || len(state.Nodes) != 0 {
		t.Fatalf("site state=%+v", state)
	}
}

func TestKubernetesClientRejectsActiveSiteWithoutWorkerNodes(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		switch {
		case request.Method == http.MethodGet && strings.Contains(request.URL.Path, "/deployments/easyai-worker-ningbo"):
			_, _ = w.Write([]byte(`{"spec":{"replicas":1,"template":{"spec":{"containers":[]}}}}`))
		case request.Method == http.MethodGet && request.URL.Path == "/api/v1/nodes":
			_, _ = w.Write([]byte(`{"items":[]}`))
		default:
			http.NotFound(w, request)
		}
	}))
	defer server.Close()
	tokenFile := filepath.Join(t.TempDir(), "token")
	if err := os.WriteFile(tokenFile, []byte("test-token"), 0o600); err != nil {
		t.Fatal(err)
	}
	client, err := NewKubernetesClient(KubernetesConfig{
		Namespace: "easyai", APIServer: server.URL, TokenFile: tokenFile, HTTPClient: server.Client(),
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := client.SiteState(context.Background(), "ningbo"); err == nil ||
		!strings.Contains(err.Error(), "site ningbo has no matching nodes") {
		t.Fatalf("expected missing Worker node error, got %v", err)
	}
}

func TestQuantityParsing(t *testing.T) {
	if value, err := parseBinaryQuantity("1536Mi"); err != nil || value != 1536<<20 {
		t.Fatalf("memory=%d err=%v", value, err)
	}
	if value, err := parseCPUQuantity("250000000n"); err != nil || value != 250 {
		t.Fatalf("cpu=%d err=%v", value, err)
	}
}