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