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

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"testing"
	"time"

	"vdi-cloud/platform/internal/model"
)

func serveRecoverySunshine(t *testing.T) {
	t.Helper()
	listener, err := net.Listen("tcp", "127.0.0.1:47989")
	if err != nil {
		t.Fatal(err)
	}
	srv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprint(w, `<root status_code="200"><hostname>test-sunshine</hostname></root>`)
	})}
	go srv.Serve(listener)
	t.Cleanup(func() { srv.Close() })
}

func seedGPURecovery(t *testing.T, f *fixture, status string) model.Task {
	t.Helper()
	task := prepareExclusive(t, f)
	state := gpuPrepareState{Phase: "start-waiting", Node: "gpu-node", Members: []int{100, 101, 102}, Terminal: "gpu-box", VMID: 102, UPID: "UPID:gpu-test"}
	data, _ := json.Marshal(state)
	if err := f.db.Get(&task, "UPDATE tasks SET status=$1,provider_ref=$2 WHERE id=$3 RETURNING *", status, string(data), task.ID); err != nil {
		t.Fatal(err)
	}
	return task
}

func TestGPURecoveryRequiresActualOwnershipAndSunshine(t *testing.T) {
	for _, scenario := range []string{"success", "original-failed", "pending-shutdown", "multiple-running", "pve-unreachable", "sunshine-unreachable", "authorization-revoked", "config-changed", "unknown-submission"} {
		t.Run(scenario, func(t *testing.T) {
			f, p := exclusiveFixture(t, "normal")
			task := seedGPURecovery(t, f, "timeout")
			p.power[100], p.power[102] = "stopped", "running"
			want := "timeout"
			switch scenario {
			case "success":
				serveRecoverySunshine(t)
				want = "success"
			case "original-failed":
				p.mode = "failed-task"
				want = "failed"
			case "pending-shutdown":
				p.power[102] = "stopped"
				want = "failed"
			case "multiple-running":
				p.power[100] = "running"
			case "pve-unreachable":
				p.mode = "poll-failure"
			case "authorization-revoked":
				if _, err := f.db.Exec("UPDATE users SET enabled=false"); err != nil {
					t.Fatal(err)
				}
				want = "failed"
			case "config-changed":
				f.server.Config.ExclusiveGPUVMIDs = []int{100, 102}
			case "unknown-submission":
				if _, err := f.db.Exec("UPDATE tasks SET provider_ref=$1 WHERE id=$2", `{"phase":"start-submitting","node":"gpu-node","members":[100,101,102],"terminal":"gpu-box","vmid":102}`, task.ID); err != nil {
					t.Fatal(err)
				}
			}
			f.server.recoverGPUPrepare(context.Background(), task)
			var result model.Task
			if err := f.db.Get(&result, "SELECT * FROM tasks WHERE id=$1", task.ID); err != nil {
				t.Fatal(err)
			}
			if result.Status != want || len(p.actions) != 0 {
				t.Fatal("incorrect read-only recovery", result.Status, want, p.actions)
			}
			if want == "success" {
				var power, service string
				f.db.QueryRow("SELECT power_state,service_state FROM desktops WHERE id=$1", task.TargetID).Scan(&power, &service)
				if power != "running" || service != "ready" {
					t.Fatal("desktop cache did not converge", power, service)
				}
				f.server.recoverGPUPrepare(context.Background(), task)
				var n int
				f.db.Get(&n, "SELECT count(*) FROM audit_logs WHERE action='auto-reconcile-gpu-switch' AND request_id=$1", task.ID)
				if n != 1 {
					t.Fatal("recovery audit was duplicated", n)
				}
			}
		})
	}
}

func TestGPUWorkerRecoversOriginalTaskAfterRestartWithoutCommands(t *testing.T) {
	f, p := exclusiveFixture(t, "normal")
	task := seedGPURecovery(t, f, "running")
	p.power[100], p.power[102] = "stopped", "running"
	serveRecoverySunshine(t)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { f.server.RunWorker(ctx); close(done) }()
	defer func() { cancel(); <-done }()
	deadline := time.Now().Add(5 * time.Second)
	for {
		var status string
		if err := f.db.Get(&status, "SELECT status FROM tasks WHERE id=$1", task.ID); err != nil {
			t.Fatal(err)
		}
		if status == "success" {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("worker did not recover orphaned task", status)
		}
		time.Sleep(20 * time.Millisecond)
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if len(p.actions) != 0 {
		t.Fatal("restart replayed a power command", p.actions)
	}
}

func TestGPUTransientReadFailureKeepsSingleCommandSequence(t *testing.T) {
	f, p := exclusiveFixture(t, "normal")
	p.readFailures = 1
	serveRecoverySunshine(t)
	task := prepareExclusive(t, f)
	f.server.executeGPUPrepare(context.Background(), task)
	var status string
	f.db.Get(&status, "SELECT status FROM tasks WHERE id=$1", task.ID)
	if status != "success" || fmt.Sprint(p.actions) != "[100:shutdown 102:start]" {
		t.Fatal("transient poll failure interrupted or replayed commands", status, p.actions)
	}
}

func TestGPURecoveryPollingRetainsTaskAndDeviceIsolation(t *testing.T) {
	f, _ := exclusiveFixture(t, "normal")
	task := seedGPURecovery(t, f, "timeout")
	code, out := f.request("GET", "/terminal/desktop-prepare/"+task.ID, "gpu-device", nil, "")
	expect(t, 200, code, out)
	if out["status"] != "running" {
		t.Fatal("terminal stopped polling recoverable task", out)
	}
	code, out = f.request("POST", "/terminal/desktop-prepare", "gpu-device", map[string]any{"desktopId": task.TargetID}, "")
	expect(t, 202, code, out)
	if out["taskId"] != task.ID {
		t.Fatal("new task submitted during recovery", out)
	}
	if _, err := f.db.Exec("UPDATE tasks SET created_at=now()-interval '6 minutes' WHERE id=$1", task.ID); err != nil {
		t.Fatal(err)
	}
	code, out = f.request("GET", "/terminal/desktop-prepare/"+task.ID, "gpu-device", nil, "")
	expect(t, 200, code, out)
	if out["status"] != "timeout" {
		t.Fatal("terminal waiting exceeded recovery window", out)
	}
}