vdi-cloud/platform/internal/server/catalog_test.go

package server

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"
	"vdi-cloud/platform/pkg/proxmox"
)

type catalogPVE struct {
	posts       atomic.Int32
	deny        atomic.Bool
	changed     atomic.Bool
	lowSpace    atomic.Bool
	pollFailure atomic.Bool
	holdPoll    atomic.Bool
	outcome     string
}

func fakeCatalogPVE(t *testing.T, f *fixture, outcome string) *catalogPVE {
	t.Helper()
	p := &catalogPVE{outcome: outcome}
	p.pollFailure.Store(outcome == "poll-failure")
	p.holdPoll.Store(outcome == "running")
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := strings.TrimPrefix(r.URL.Path, "/api2/json")
		var data any
		switch {
		case path == "/cluster/resources":
			data = []map[string]any{{"type": "qemu", "node": "node-a", "vmid": 8000, "template": 1, "name": "ubuntu-template", "status": "stopped"}}
		case path == "/nodes/node-a/qemu/8000/config":
			digest := "v1"
			if p.changed.Load() {
				digest = "v2"
			}
			data = map[string]any{"template": 1, "digest": digest, "scsi0": "local:base-8000,size=4G"}
		case path == "/nodes/node-a/qemu/8000/feature":
			data = map[string]any{"hasFeature": 1}
		case path == "/access/permissions":
			allowed := 1
			if p.deny.Load() {
				allowed = 0
			}
			data = map[string]any{r.URL.Query().Get("path"): map[string]int{"VM.Clone": allowed, "VM.Allocate": allowed, "Datastore.AllocateSpace": allowed}}
		case path == "/nodes/node-a/storage":
			avail := float64(100 * (1 << 30))
			if p.lowSpace.Load() {
				avail = 0
			}
			data = []map[string]any{{"storage": "local", "avail": avail, "active": 1, "enabled": 1, "content": "images"}}
		case path == "/nodes/node-a/qemu/8000/clone":
			p.posts.Add(1)
			if err := r.ParseForm(); err != nil {
				t.Error(err)
			}
			if r.Method != "POST" || r.Form.Get("name") != "vdi-9000" || r.Form.Get("newid") != "9000" || r.Form.Get("full") != "1" || r.Form.Get("storage") != "local" {
				t.Errorf("incorrect clone form: %v", r.Form)
			}
			if outcome == "ambiguous" {
				w.WriteHeader(500)
				return
			}
			data = "UPID:test-clone"
		case strings.Contains(path, "/tasks/"):
			if p.pollFailure.Load() {
				w.WriteHeader(503)
				return
			}
			if p.holdPoll.Load() {
				data = map[string]any{"status": "running"}
				break
			}
			exit := "OK"
			if outcome == "failed" {
				exit = "clone failed"
			}
			data = map[string]any{"status": "stopped", "exitstatus": exit}
		case path == "/nodes/node-a/qemu/9000/config":
			data = map[string]any{"name": "vdi-9000", "scsi0": "local:vm-9000,size=4G"}
		case path == "/nodes/node-a/qemu/9000/status/current":
			data = map[string]any{"status": "stopped"}
		default:
			t.Errorf("unexpected PVE request: %s", path)
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
	}))
	t.Cleanup(srv.Close)
	f.server.PVE = &proxmox.Client{BaseURL: srv.URL, HTTP: srv.Client()}
	return p
}
func registerCatalog(t *testing.T, f *fixture, admin string) (string, string) {
	t.Helper()
	status, out := f.request("POST", "/templates", admin, map[string]any{"name": "Ubuntu base", "version": "1.0", "os": "Ubuntu", "node": "node-a", "vmid": 8000}, "")
	expect(t, 201, status, out)
	tpl := out["id"].(string)
	status, out = f.request("POST", "/pools", admin, map[string]any{"name": "Engineering", "type": "dedicated", "maxCount": 1, "templateId": tpl, "enabled": true}, "")
	expect(t, 200, status, out)
	return tpl, out["id"].(string)
}
func cloneBody(tpl, pool string) map[string]any {
	return map[string]any{"templateId": tpl, "poolId": pool, "name": "New desktop", "vmid": 9000, "full": true, "storage": "local"}
}
func waitCatalogTask(t *testing.T, f *fixture, admin, id string) map[string]any {
	t.Helper()
	deadline := time.Now().Add(6 * time.Second)
	for {
		status, out := f.request("GET", "/tasks/"+id, admin, nil, "")
		expect(t, 200, status, out)
		if out["status"] != "queued" && out["status"] != "running" {
			return out
		}
		if time.Now().After(deadline) {
			t.Fatal("clone never completed", out)
		}
		time.Sleep(30 * time.Millisecond)
	}
}
func TestCloneCreatesDurableUninitializedDesktop(t *testing.T) {
	f := newFixture(t)
	p := fakeCatalogPVE(t, f, "success")
	admin := f.login("admin")
	tpl, pool := registerCatalog(t, f, admin)
	body := cloneBody(tpl, pool)
	status, out := f.request("POST", "/desktops/clone/check", admin, body, "")
	expect(t, 200, status, out)
	if out["available"] != true {
		t.Fatal(out)
	}
	status, out = f.request("POST", "/desktops/clone", admin, body, "clone-once")
	expect(t, 202, status, out)
	id := out["id"].(string)
	task := waitCatalogTask(t, f, admin, id)
	if task["status"] != "success" {
		t.Fatal(task)
	}
	status, again := f.request("POST", "/desktops/clone", admin, body, "clone-once")
	expect(t, 202, status, again)
	if again["id"] != id || p.posts.Load() != 1 {
		t.Fatal("clone replayed", again, p.posts.Load())
	}
	body["name"] = "different"
	status, out = f.request("POST", "/desktops/clone", admin, body, "clone-once")
	expect(t, 409, status, out)
	status, out = f.request("GET", "/desktops/"+task["targetId"].(string), admin, nil, "")
	expect(t, 200, status, out)
	if out["setupRequired"] != true || out["endpoint"] != "" || out["poolId"] != pool {
		t.Fatal("new VM incorrectly advertised ready", out)
	}
	connection := map[string]any{"endpoint": "192.0.2.20", "width": 1920, "height": 1080, "fps": 30, "initialized": false}
	desk := task["targetId"].(string)
	status, out = f.request("PUT", "/desktops/"+desk+"/connection", admin, connection, "")
	expect(t, 400, status, out)
	connection["initialized"] = true
	connection["endpoint"] = "https://invalid"
	status, out = f.request("PUT", "/desktops/"+desk+"/connection", admin, connection, "")
	expect(t, 400, status, out)
	connection["endpoint"] = "192.0.2.20"
	status, out = f.request("PUT", "/desktops/"+desk+"/connection", admin, connection, "")
	expect(t, 200, status, out)
	if out["setupRequired"] != false {
		t.Fatal(out)
	}
	body["vmid"] = 9001
	status, out = f.request("POST", "/desktops/clone/check", admin, body, "")
	expect(t, 200, status, out)
	if out["available"] != false {
		t.Fatal("full pool passed preflight", out)
	}
	status, out = f.request("POST", "/desktops/clone", admin, body, "over-capacity")
	expect(t, 409, status, out)
	var count int
	_ = f.db.Get(&count, "SELECT count(*) FROM audit_logs WHERE action='clone' AND result='success'")
	if count != 1 {
		t.Fatal("missing final audit")
	}
}
func TestClonePreflightAndCatalogRBAC(t *testing.T) {
	f := newFixture(t)
	p := fakeCatalogPVE(t, f, "success")
	admin := f.login("admin")
	tpl, pool := registerCatalog(t, f, admin)
	body := cloneBody(tpl, pool)
	for _, role := range []string{"viewer", "operator"} {
		token := f.login(role)
		for _, endpoint := range []string{"/templates", "/pools", "/desktops/clone", "/desktops/clone/check"} {
			status, out := f.request("POST", endpoint, token, body, "forbidden")
			expect(t, 403, status, out)
		}
	}
	for _, tc := range []struct {
		name string
		flag *atomic.Bool
	}{{"permission", &p.deny}, {"template-changed", &p.changed}, {"storage-capacity", &p.lowSpace}} {
		t.Run(tc.name, func(t *testing.T) {
			tc.flag.Store(true)
			defer tc.flag.Store(false)
			status, out := f.request("POST", "/desktops/clone/check", admin, body, "")
			expect(t, 200, status, out)
			if out["available"] != false {
				t.Fatal(out)
			}
			status, out = f.request("POST", "/desktops/clone", admin, body, tc.name)
			expect(t, 409, status, out)
		})
	}
	if p.posts.Load() != 0 {
		t.Fatal("preflight failure still issued clone")
	}
	status, out := f.request("PATCH", "/templates/"+tpl, admin, map[string]any{"enabled": false}, "")
	expect(t, 200, status, out)
	status, out = f.request("POST", "/desktops/clone", admin, body, "disabled")
	expect(t, 409, status, out)
}
func TestAmbiguousAndFailedCloneRetainPoolCapacity(t *testing.T) {
	for _, outcome := range []string{"ambiguous", "failed"} {
		t.Run(outcome, func(t *testing.T) {
			f := newFixture(t)
			p := fakeCatalogPVE(t, f, outcome)
			admin := f.login("admin")
			tpl, pool := registerCatalog(t, f, admin)
			status, out := f.request("POST", "/desktops/clone", admin, cloneBody(tpl, pool), "uncertain")
			expect(t, 202, status, out)
			task := waitCatalogTask(t, f, admin, out["id"].(string))
			want := "timeout"
			if outcome == "failed" {
				want = "failed"
			}
			if task["status"] != want {
				t.Fatal(task)
			}
			var state string
			_ = f.db.Get(&state, "SELECT state FROM provisioning_jobs WHERE task_id=$1", task["id"])
			if state != "uncertain" {
				t.Fatal("lost uncertain reservation", state)
			}
			status, out = f.request("GET", "/pools", admin, nil, "")
			expect(t, 200, status, out)
			if out["items"].([]any)[0].(map[string]any)["reservedCount"] != float64(1) {
				t.Fatal(out)
			}
			time.Sleep(350 * time.Millisecond)
			if p.posts.Load() != 1 {
				t.Fatal("ambiguous POST replayed")
			}
		})
	}
}
func TestKnownUPIDReconcilesWithoutAnotherClone(t *testing.T) {
	f := newFixture(t)
	p := fakeCatalogPVE(t, f, "poll-failure")
	admin := f.login("admin")
	tpl, pool := registerCatalog(t, f, admin)
	status, out := f.request("POST", "/desktops/clone", admin, cloneBody(tpl, pool), "resume-known")
	expect(t, 202, status, out)
	id := out["id"].(string)
	task := waitCatalogTask(t, f, admin, id)
	if task["status"] != "timeout" || task["providerRef"] == "" {
		t.Fatal(task)
	}
	p.pollFailure.Store(false)
	status, out = f.request("POST", "/tasks/"+id+"/reconcile", admin, nil, "")
	expect(t, 202, status, out)
	task = waitCatalogTask(t, f, admin, id)
	if task["status"] != "success" || p.posts.Load() != 1 {
		t.Fatal("known UPID was not safely resumed", task, p.posts.Load())
	}
	status, out = f.request("POST", "/tasks/"+id+"/reconcile", admin, nil, "")
	expect(t, 409, status, out)
}

func TestCloneDoesNotBlockOrdinaryTasks(t *testing.T) {
	f := newFixture(t)
	p := fakeCatalogPVE(t, f, "running")
	admin := f.login("admin")
	tpl, pool := registerCatalog(t, f, admin)
	status, out := f.request("POST", "/desktops/clone", admin, cloneBody(tpl, pool), "long-clone")
	expect(t, 202, status, out)
	cloneID := out["id"].(string)
	deadline := time.Now().Add(5 * time.Second)
	for p.posts.Load() == 0 {
		if time.Now().After(deadline) {
			t.Fatal("clone did not start")
		}
		time.Sleep(20 * time.Millisecond)
	}
	status, out = f.request("POST", "/desktops/desk_demo_kylin/actions", admin, map[string]any{"action": "start"}, "during-clone")
	expect(t, 202, status, out)
	power := waitCatalogTask(t, f, admin, out["id"].(string))
	if power["status"] != "success" {
		t.Fatal("ordinary task blocked by clone", power)
	}
	status, out = f.request("GET", "/tasks/"+cloneID, admin, nil, "")
	expect(t, 200, status, out)
	if out["status"] != "running" {
		t.Fatal("clone should still be polling", out)
	}
	p.holdPoll.Store(false)
	if task := waitCatalogTask(t, f, admin, cloneID); task["status"] != "success" {
		t.Fatal(task)
	}
}

func TestPoolMembershipCapacityAndConnectionBusy(t *testing.T) {
	f := newFixture(t)
	admin := f.login("admin")
	input := map[string]any{"name": "Pool A", "type": "shared", "maxCount": 1, "templateId": nil, "enabled": true}
	status, out := f.request("POST", "/pools", admin, input, "")
	expect(t, 200, status, out)
	pool := out["id"].(string)
	status, out = f.request("PUT", "/pools/"+pool+"/members", admin, map[string]any{"desktopIds": []string{"desk_demo_ubuntu", "desk_demo_kylin"}}, "")
	expect(t, 409, status, out)
	status, out = f.request("PUT", "/pools/"+pool+"/members", admin, map[string]any{"desktopIds": []string{"desk_demo_ubuntu"}}, "")
	expect(t, 200, status, out)
	input["name"] = "Pool B"
	status, out = f.request("POST", "/pools", admin, input, "")
	expect(t, 200, status, out)
	other := out["id"].(string)
	status, out = f.request("PUT", "/pools/"+other+"/members", admin, map[string]any{"desktopIds": []string{"desk_demo_ubuntu"}}, "")
	expect(t, 409, status, out)
	input["type"] = "dedicated"
	input["name"] = "Renamed"
	status, out = f.request("PUT", "/pools/"+pool, admin, input, "")
	expect(t, 409, status, out)
	input["type"] = "shared"
	status, out = f.request("PUT", "/pools/"+pool, admin, input, "")
	expect(t, 200, status, out)
	var name string
	_ = f.db.Get(&name, "SELECT pool FROM desktops WHERE id='desk_demo_ubuntu'")
	if name != "Renamed" {
		t.Fatal("pool display name stale", name)
	}
	_, err := f.db.Exec("UPDATE terminals SET session_desktop_id='desk_demo_ubuntu',last_seen=now() WHERE id='term_demo_001'")
	if err != nil {
		t.Fatal(err)
	}
	status, out = f.request("PUT", "/desktops/desk_demo_ubuntu/connection", admin, map[string]any{"endpoint": "192.0.2.1", "width": 1920, "height": 1080, "fps": 30, "initialized": true}, "")
	expect(t, 409, status, out)
}