Agent-System/internal/taskrun/store_lease_pg_test.go

package taskrun_test

import (
	"context"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/db"
	"chat-agent-system/internal/taskrun"
)

func openLeaseTestPool(t *testing.T) *pgxpool.Pool {
	t.Helper()
	dsn := db.SkipIfNoTestDatabaseURL(t)
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		t.Skipf("postgres unavailable: %v", err)
		return nil
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		t.Skipf("postgres ping failed: %v", err)
		return nil
	}
	migCtx, migCancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer migCancel()
	if err := db.RunMigrations(migCtx, pool, db.ResolveMigrationsDir()); err != nil {
		pool.Close()
		t.Skipf("migration failed: %v", err)
		return nil
	}
	if !taskrun.LeaseColumnsReady(ctx, pool) {
		pool.Close()
		t.Skip("task run lease columns not ready")
	}
	return pool
}

func TestPGStore_StaleLeaseMarkedInterrupted(t *testing.T) {
	pool := openLeaseTestPool(t)
	if pool == nil {
		return
	}
	defer pool.Close()

	ctx := context.Background()
	store := taskrun.NewPGStore(pool)
	runID := "run_stale_lease_" + time.Now().Format("150405.000000")

	var userID int64
	err := pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID)
	if err != nil {
		t.Fatalf("fetch user: %v", err)
	}

	expired := time.Now().Add(-1 * time.Minute)
	_, err = pool.Exec(ctx, `
		INSERT INTO task_runs (
			run_id, session_id, user_id, status, worker_id, lease_expires_at, lease_heartbeat_at, created_at, updated_at
		) VALUES ($1, 'sess', $2, 'running', 'dead-worker', $3, $3, now(), now())
	`, runID, userID, expired)
	if err != nil {
		t.Fatalf("insert stale run: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})

	n, err := store.MarkOrphanedRunning(ctx)
	if err != nil {
		t.Fatalf("MarkOrphanedRunning: %v", err)
	}
	if n < 1 {
		t.Fatalf("expected at least 1 interrupted, got %d", n)
	}

	rec, err := store.GetRun(ctx, runID)
	if err != nil {
		t.Fatalf("GetRun: %v", err)
	}
	if rec.Status != "interrupted" {
		t.Fatalf("status=%s want interrupted", rec.Status)
	}
}

func TestPGStore_FreshLeaseSurvivesOrphanScan(t *testing.T) {
	pool := openLeaseTestPool(t)
	if pool == nil {
		return
	}
	defer pool.Close()

	ctx := context.Background()
	store := taskrun.NewPGStore(pool)
	runID := "run_fresh_lease_" + time.Now().Format("150405.000000")
	workerID := "worker-alive"

	var userID int64
	_ = pool.QueryRow(ctx, `SELECT id FROM users ORDER BY id LIMIT 1`).Scan(&userID)

	fresh := time.Now().Add(30 * time.Second)
	_, err := pool.Exec(ctx, `
		INSERT INTO task_runs (
			run_id, session_id, user_id, status, worker_id, lease_expires_at, lease_heartbeat_at, created_at, updated_at
		) VALUES ($1, 'sess', $2, 'running', $3, $4, now(), now(), now())
	`, runID, userID, workerID, fresh)
	if err != nil {
		t.Fatalf("insert fresh run: %v", err)
	}
	t.Cleanup(func() {
		_, _ = pool.Exec(context.Background(), `DELETE FROM task_runs WHERE run_id = $1`, runID)
	})

	if _, err := store.MarkOrphanedRunning(ctx); err != nil {
		t.Fatalf("MarkOrphanedRunning: %v", err)
	}

	rec, err := store.GetRun(ctx, runID)
	if err != nil {
		t.Fatalf("GetRun: %v", err)
	}
	if rec.Status != "running" {
		t.Fatalf("fresh lease run should stay running, got %s", rec.Status)
	}
}