geo/server/internal/tenant/app/desktop_task_service_test.go

package app

import (
	"regexp"
	"strings"
	"testing"
)

func TestDesktopTaskRecoverySelectQueryMatchesRecoveredLeaseScanOrder(t *testing.T) {
	t.Parallel()

	query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeStartup)
	gotColumns := recoverDesktopTaskSelectColumns(query)
	wantColumns := []string{
		"desktop_id",
		"workspace_id",
		"kind",
		"status",
		"active_attempt_id",
		"attempts",
		"publish_submit_started_at",
	}

	if len(gotColumns) != len(wantColumns) {
		t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns)
	}
	for index := range wantColumns {
		if gotColumns[index] != wantColumns[index] {
			t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns)
		}
	}
}

func TestDesktopTaskRecoverySelectQueryLeaseExpiryFiltersExpiredLeases(t *testing.T) {
	t.Parallel()

	query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeLeaseExpiry)

	for _, fragment := range []string{
		"AND lease_expires_at IS NOT NULL",
		"AND lease_expires_at < NOW()",
		"FOR UPDATE",
	} {
		if !strings.Contains(query, fragment) {
			t.Fatalf("recover lease-expiry query missing %q: %s", fragment, query)
		}
	}
}

func TestReconcileRetryPublishesTaskAvailableForAllDesktopTaskKinds(t *testing.T) {
	t.Parallel()

	for _, kind := range []string{"publish", "monitor"} {
		if got := reconcileDesktopTaskEventType(kind, "retry"); got != "task_available" {
			t.Fatalf("reconcileDesktopTaskEventType(%q, retry) = %q, want task_available", kind, got)
		}
	}
}

func TestDesktopTaskRecoveryPayloadPublishLeaseExpiryRequeuesInsteadOfUnknown(t *testing.T) {
	t.Parallel()

	payload, reason, message, err := desktopTaskRecoveryPayload(desktopTaskRecoveryModeLeaseExpiry, "publish")
	if err != nil {
		t.Fatalf("desktopTaskRecoveryPayload returned error: %v", err)
	}
	if reason != "lease_expired" {
		t.Fatalf("reason = %q, want lease_expired", reason)
	}
	if !strings.Contains(message, "re-queued for retry") {
		t.Fatalf("message = %q, want re-queued for retry", message)
	}
	if strings.Contains(message, "unknown") {
		t.Fatalf("message must not move publish recovery to unknown: %q", message)
	}
	if !strings.Contains(string(payload), "desktop_task_lease_expiry") {
		t.Fatalf("payload missing source: %s", payload)
	}
}

func TestDesktopTaskRecoveryPayloadPublishFinalFailsAfterMaxAttempts(t *testing.T) {
	t.Parallel()

	payload, _, message, err := desktopTaskRecoveryPayload(desktopTaskRecoveryModeLeaseExpiry, "publish_final")
	if err != nil {
		t.Fatalf("desktopTaskRecoveryPayload returned error: %v", err)
	}
	if !strings.Contains(message, "marked failed") {
		t.Fatalf("message = %q, want marked failed", message)
	}
	if !strings.Contains(message, "3 attempts") {
		t.Fatalf("message = %q, want max attempts", message)
	}
	if !strings.Contains(string(payload), "marked failed") {
		t.Fatalf("payload missing final failure message: %s", payload)
	}
}

func TestReconcileNonRetryPublishesTaskReconciled(t *testing.T) {
	t.Parallel()

	if got := reconcileDesktopTaskEventType("publish", "failed"); got != "task_reconciled" {
		t.Fatalf("reconcileDesktopTaskEventType(publish, failed) = %q, want task_reconciled", got)
	}
}

func TestResolvePublishRecoveryOutcome(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name          string
		submitStarted bool
		attempts      int
		wantStatus    string
		wantKind      string
	}{
		{"pre-submit first attempt requeues", false, 1, "queued", "publish"},
		{"pre-submit just below cap requeues", false, desktopPublishMaxAttempts - 1, "queued", "publish"},
		{"pre-submit at cap fails terminally", false, desktopPublishMaxAttempts, "failed", "publish_final"},
		{"pre-submit above cap fails terminally", false, desktopPublishMaxAttempts + 1, "failed", "publish_final"},
		// Submit may have happened: never auto-requeue or silently fail — route to unknown
		// (manual reconcile) regardless of remaining attempts, to avoid duplicate publishing.
		{"submit started first attempt is uncertain", true, 1, "unknown", "publish_uncertain"},
		{"submit started at cap is still uncertain", true, desktopPublishMaxAttempts, "unknown", "publish_uncertain"},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			gotStatus, gotKind := resolvePublishRecoveryOutcome(tc.submitStarted, tc.attempts)
			if gotStatus != tc.wantStatus || gotKind != tc.wantKind {
				t.Fatalf("resolvePublishRecoveryOutcome(%v, %d) = (%q, %q), want (%q, %q)",
					tc.submitStarted, tc.attempts, gotStatus, gotKind, tc.wantStatus, tc.wantKind)
			}
		})
	}
}

func recoverDesktopTaskSelectColumns(query string) []string {
	re := regexp.MustCompile(`(?is)select\s+(.*?)\s+from\s+desktop_tasks`)
	match := re.FindStringSubmatch(query)
	if len(match) != 2 {
		return nil
	}

	rawColumns := strings.Split(match[1], ",")
	columns := make([]string, 0, len(rawColumns))
	for _, rawColumn := range rawColumns {
		column := strings.TrimSpace(rawColumn)
		column = strings.Fields(column)[0]
		columns = append(columns, column)
	}
	return columns
}