proxui/worker/internal/tasks/proxmox_poll_test.go

package tasks

import (
	"context"
	"log/slog"
	"testing"
	"time"

	"github.com/hibiken/asynq"
)

func TestProxmoxTaskPollHandlerMarksVMRunningOnOK(t *testing.T) {
	resolver := &stubTaskClientResolver{client: &stubTaskClient{
		statuses: []ProxmoxTaskStatus{
			{Status: "running"},
			{Status: "stopped", ExitStatus: "OK"},
		},
	}}
	vmStore := &stubVMStatusStore{}
	auditWriter := &stubAuditWriter{}
	handler := NewProxmoxTaskPollHandler(resolver, vmStore, auditWriter, slog.Default()).WithPolling(3, 0)

	task := newPollTask(t, ProxmoxTaskPollPayload{
		ClusterID:     "cluster-1",
		Node:          "pve",
		UPID:          "UPID:pve:1",
		TargetType:    "vm",
		TargetID:      "vm-1",
		TenantID:      "tenant-1",
		ProfileID:     "profile-1",
		Action:        "vm.power.start",
		SuccessStatus: "running",
	})

	if err := handler.ProcessTask(context.Background(), task); err != nil {
		t.Fatalf("ProcessTask() error = %v", err)
	}

	if resolver.clusterID != "cluster-1" {
		t.Fatalf("clusterID = %q, want cluster-1", resolver.clusterID)
	}
	if vmStore.status != "running" {
		t.Fatalf("vm status = %q, want running", vmStore.status)
	}
	if auditWriter.event.Action != "vm.power.start" {
		t.Fatalf("audit action = %q, want vm.power.start", auditWriter.event.Action)
	}
	if auditWriter.event.Metadata["exit_status"] != "OK" {
		t.Fatalf("audit exit_status = %v, want OK", auditWriter.event.Metadata["exit_status"])
	}
}

func TestProxmoxTaskPollHandlerMarksVMFailedOnErrorExit(t *testing.T) {
	resolver := &stubTaskClientResolver{client: &stubTaskClient{
		statuses: []ProxmoxTaskStatus{
			{Status: "stopped", ExitStatus: "error"},
		},
	}}
	vmStore := &stubVMStatusStore{}
	auditWriter := &stubAuditWriter{}
	handler := NewProxmoxTaskPollHandler(resolver, vmStore, auditWriter, slog.Default()).WithPolling(2, 0)

	task := newPollTask(t, validPollPayload())

	if err := handler.ProcessTask(context.Background(), task); err != nil {
		t.Fatalf("ProcessTask() error = %v", err)
	}

	if vmStore.status != "failed" {
		t.Fatalf("vm status = %q, want failed", vmStore.status)
	}
	if auditWriter.event.Metadata["target_status"] != "failed" {
		t.Fatalf("audit target_status = %v, want failed", auditWriter.event.Metadata["target_status"])
	}
}

func TestProxmoxTaskPollHandlerReturnsErrorWhenStillRunning(t *testing.T) {
	resolver := &stubTaskClientResolver{client: &stubTaskClient{
		statuses: []ProxmoxTaskStatus{
			{Status: "running"},
			{Status: "running"},
		},
	}}
	handler := NewProxmoxTaskPollHandler(resolver, &stubVMStatusStore{}, &stubAuditWriter{}, slog.Default()).WithPolling(2, 0)

	err := handler.ProcessTask(context.Background(), newPollTask(t, validPollPayload()))
	if err == nil {
		t.Fatal("ProcessTask() error = nil, want error")
	}
}

func TestProxmoxTaskPollHandlerRejectsInvalidPayload(t *testing.T) {
	handler := NewProxmoxTaskPollHandler(&stubTaskClientResolver{}, &stubVMStatusStore{}, &stubAuditWriter{}, slog.Default())

	err := handler.ProcessTask(context.Background(), asynq.NewTask(TypeProxmoxTaskPoll, []byte(`{"cluster_id":""}`)))
	if err == nil {
		t.Fatal("ProcessTask() error = nil, want error")
	}
}

func newPollTask(t *testing.T, payload ProxmoxTaskPollPayload) *asynq.Task {
	t.Helper()

	task, err := NewProxmoxTaskPollTask(payload)
	if err != nil {
		t.Fatalf("NewProxmoxTaskPollTask() error = %v", err)
	}
	return task
}

func validPollPayload() ProxmoxTaskPollPayload {
	return ProxmoxTaskPollPayload{
		ClusterID:     "cluster-1",
		Node:          "pve",
		UPID:          "UPID:pve:1",
		TargetType:    "vm",
		TargetID:      "vm-1",
		TenantID:      "tenant-1",
		Action:        "vm.power.reboot",
		SuccessStatus: "running",
	}
}

type stubTaskClientResolver struct {
	clusterID string
	client    ProxmoxTaskClient
}

func (s *stubTaskClientResolver) ResolveTaskClient(_ context.Context, clusterID string) (ProxmoxTaskClient, error) {
	s.clusterID = clusterID
	return s.client, nil
}

type stubTaskClient struct {
	statuses []ProxmoxTaskStatus
	calls    int
}

func (s *stubTaskClient) GetTaskStatus(_ context.Context, node string, upid string) (ProxmoxTaskStatus, error) {
	s.calls++
	if node != "pve" {
		return ProxmoxTaskStatus{}, nil
	}
	if upid == "" {
		return ProxmoxTaskStatus{}, nil
	}
	if len(s.statuses) == 0 {
		return ProxmoxTaskStatus{Status: "running"}, nil
	}

	status := s.statuses[0]
	s.statuses = s.statuses[1:]
	return status, nil
}

type stubVMStatusStore struct {
	vmID   string
	status string
}

func (s *stubVMStatusStore) SetVMStatus(_ context.Context, vmID string, status string) error {
	s.vmID = vmID
	s.status = status
	return nil
}

type stubAuditWriter struct {
	event AuditEvent
}

func (s *stubAuditWriter) WriteAudit(_ context.Context, event AuditEvent) error {
	s.event = event
	return nil
}

func TestSleepContextReturnsContextError(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	if err := sleepContext(ctx, time.Second); err == nil {
		t.Fatal("sleepContext() error = nil, want context error")
	}
}