queuer_test.go

package queuer

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

	"github.com/google/uuid"
	"github.com/siherrmann/queuer/helper"
	"github.com/siherrmann/queuer/model"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestNewQueuer(t *testing.T) {
	tests := []struct {
		name           string
		maxConcurrency int
		options        []*model.OnError
		dbEnvs         map[string]string
		expectError    bool
	}{
		{
			name:           "Valid queuer",
			maxConcurrency: 100,
			options:        nil,
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST":     "localhost",
				"QUEUER_DB_PORT":     dbPort,
				"QUEUER_DB_DATABASE": "database",
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: false,
		},
		{
			name:           "Valid queuer with options",
			maxConcurrency: 100,
			options: []*model.OnError{
				{
					Timeout:      10.0,
					MaxRetries:   3,
					RetryDelay:   1.0,
					RetryBackoff: model.RETRY_BACKOFF_LINEAR,
				},
			},
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST":     "localhost",
				"QUEUER_DB_PORT":     dbPort,
				"QUEUER_DB_DATABASE": "database",
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: false,
		},
		{
			name:           "Invalid max concurrency",
			maxConcurrency: -1,
			options:        nil,
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST":     "localhost",
				"QUEUER_DB_PORT":     dbPort,
				"QUEUER_DB_DATABASE": "database",
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: true,
		},
		{
			name:           "Invalid options",
			maxConcurrency: 100,
			options: []*model.OnError{
				{
					Timeout:      -10.0, // Invalid timeout value
					MaxRetries:   3,
					RetryDelay:   1.0,
					RetryBackoff: model.RETRY_BACKOFF_LINEAR,
				},
			},
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST":     "localhost",
				"QUEUER_DB_PORT":     dbPort,
				"QUEUER_DB_DATABASE": "database",
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: true,
		},
		{
			name:           "Missing DB environment variable",
			maxConcurrency: 100,
			options:        nil,
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST": "localhost",
				"QUEUER_DB_PORT": dbPort,
				// "QUEUER_DB_DATABASE": "database", // Intentionally missing
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
			},
			expectError: true,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			for key, value := range test.dbEnvs {
				t.Setenv(key, value)
			}

			if test.expectError {
				defer func() {
					if r := recover(); r == nil {
						t.Errorf("Expected panic for %s, but did not get one", test.name)
					}
				}()
			}

			queuer := NewQueuer(test.name, test.maxConcurrency, test.options...)
			if !test.expectError {
				require.NotNil(t, queuer, "Expected Queuer to be created successfully")
				assert.Equal(t, test.name, queuer.worker.Name, "Expected Queuer name to match")
				assert.Equal(t, test.maxConcurrency, queuer.worker.MaxConcurrency, "Expected Queuer max concurrency to match")
			}
		})
	}
}

func TestNewStaticQueuer(t *testing.T) {
	tests := []struct {
		name        string
		logLevel    slog.Leveler
		dbConfig    *helper.DatabaseConfiguration
		dbEnvs      map[string]string
		expectError bool
	}{
		{
			name:     "Valid static queuer with nil dbConfig",
			logLevel: slog.LevelInfo,
			dbConfig: nil,
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST":     "localhost",
				"QUEUER_DB_PORT":     dbPort,
				"QUEUER_DB_DATABASE": "database",
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: false,
		},
		{
			name:     "Valid static queuer with provided dbConfig",
			logLevel: slog.LevelInfo,
			dbConfig: &helper.DatabaseConfiguration{
				Host:          "localhost",
				Port:          dbPort,
				Database:      "database",
				Username:      "user",
				Password:      "password",
				Schema:        "public",
				SSLMode:       "disable",
				WithTableDrop: true,
			},
			dbEnvs:      map[string]string{}, // No env vars needed when dbConfig is provided
			expectError: false,
		},
		{
			name:     "Missing DB environment variable when dbConfig is nil",
			logLevel: slog.LevelInfo,
			dbConfig: nil,
			dbEnvs: map[string]string{
				"QUEUER_DB_HOST": "localhost",
				"QUEUER_DB_PORT": dbPort,
				// "QUEUER_DB_DATABASE": "database", // Intentionally missing
				"QUEUER_DB_USERNAME": "user",
				"QUEUER_DB_PASSWORD": "password",
				"QUEUER_DB_SCHEMA":   "public",
				"QUEUER_DB_SSLMODE":  "disable",
			},
			expectError: true,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			for key, value := range test.dbEnvs {
				t.Setenv(key, value)
			}

			if test.expectError {
				defer func() {
					if r := recover(); r == nil {
						t.Errorf("Expected panic for %s, but did not get one", test.name)
					}
				}()
			}

			queuer := NewStaticQueuer(test.logLevel, test.dbConfig)
			if !test.expectError {
				require.NotNil(t, queuer, "Expected StaticQueuer to be created successfully")
				assert.NotNil(t, queuer.log, "Expected logger to be initialized")
				assert.NotNil(t, queuer.DB, "Expected database connection to be initialized")
				assert.NotNil(t, queuer.dbJob, "Expected job database handler to be initialized")
				assert.NotNil(t, queuer.dbWorker, "Expected worker database handler to be initialized")
				assert.NotNil(t, queuer.dbMaster, "Expected master database handler to be initialized")
				assert.NotNil(t, queuer.tasks, "Expected tasks map to be initialized")
				assert.NotNil(t, queuer.nextIntervalFuncs, "Expected nextIntervalFuncs map to be initialized")
				assert.Nil(t, queuer.worker, "Expected worker to be nil in StaticQueuer")
			}
		})
	}
}

func TestStart(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	t.Run("Start Queuer with valid context", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		startQueuer := func() {
			queuer.Start(ctx, cancel)
		}
		assert.NotPanics(t, startQueuer, "Expected no panic when starting Queuer with valid context")
		queuer.Stop()
	})

	t.Run("Start Queuer with nil context", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")

		startQueuer := func() {
			queuer.Start(context.TODO(), nil)
		}
		assert.Panics(t, startQueuer, "Expected panic when starting Queuer with nil context")
	})

	t.Run("Start Queuer with MasterSettings", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		masterSettings := &model.MasterSettings{
			MasterPollInterval: 5 * time.Second,
			JobDeleteThreshold: 30 * 24 * time.Hour, // 30 days
		}
		masterSettings.SetDefault() // Set default values for comparison
		queuer.Start(ctx, cancel, masterSettings)

		time.Sleep(6 * time.Second) // Wait for the ticker to run at least once

		master, err := queuer.dbMaster.SelectMaster()
		assert.NoError(t, err, "Expected no error selecting master")
		require.NotNil(t, master, "Expected master to not be nil")
		assert.Equal(t, queuer.worker.RID, master.WorkerRID, "Expected master RID to match worker RID")
		assert.Equal(t, queuer.worker.ID, master.WorkerID, "Expected master ID to match worker ID")
		assert.Equal(t, queuer.worker.Status, model.WorkerStatusRunning, "Expected worker status to be RUNNING")
		assert.Equal(t, masterSettings.JobDeleteThreshold, master.Settings.JobDeleteThreshold, "Expected master retention archive to match")

		// Check if the master updated_at is within the last 2 seconds,
		// as the ticker runs every 5 seconds and we waited 6 seconds.
		assert.GreaterOrEqual(t, master.UpdatedAt.Unix(), time.Now().Add(-2*time.Second).Unix(), "Expected master updated_at to be less than or equal to current time")
	})

	t.Run("Start 2 Queuers with MasterSettings and cancel first", func(t *testing.T) {
		queuer1 := NewQueuer("test", 10)
		require.NotNil(t, queuer1, "Expected Queuer 1 to be created successfully")
		ctx1, cancel1 := context.WithCancel(context.Background())
		defer cancel1()
		masterSettings1 := &model.MasterSettings{
			MasterLockTimeout:  10 * time.Second, // Short timeout for test
			MasterPollInterval: 5 * time.Second,
			JobDeleteThreshold: 30 * 24 * time.Hour,
		}
		masterSettings1.SetDefault() // Set default values for comparison
		queuer1.Start(ctx1, cancel1, masterSettings1)

		queuer2 := NewQueuer("test", 20)
		require.NotNil(t, queuer2, "Expected Queuer 2 to be created successfully")
		ctx2, cancel2 := context.WithCancel(context.Background())
		defer cancel2()
		masterSettings2 := &model.MasterSettings{
			MasterLockTimeout:  10 * time.Second, // Short timeout for test
			MasterPollInterval: 3 * time.Second,
			JobDeleteThreshold: 20 * 24 * time.Hour,
		}
		masterSettings2.SetDefault() // Set default values for comparison
		queuer2.Start(ctx2, cancel2, masterSettings2)

		time.Sleep(6 * time.Second) // Wait for the ticker to run at least once

		master, err := queuer1.dbMaster.SelectMaster()
		assert.NoError(t, err, "Expected no error selecting master")
		require.NotNil(t, master, "Expected master to not be nil")
		assert.Equal(t, queuer1.worker.RID, master.WorkerRID, "Expected master RID to match worker RID")
		assert.Equal(t, queuer1.worker.ID, master.WorkerID, "Expected master ID to match worker ID")
		assert.Equal(t, *masterSettings1, master.Settings, "Expected master settings to match")

		// Check if the master updated_at is within the last 5 seconds,
		// as the ticker runs every 5 seconds and we waited 6 seconds.
		assert.GreaterOrEqual(t, master.UpdatedAt.Unix(), time.Now().Add(-5*time.Second).Unix(), "Expected master updated_at to be in the last 5 seconds")

		// Cancel the first queuer
		cancel1()

		time.Sleep(12 * time.Second) // Wait for master lock timeout (10s) + buffer

		master, err = queuer2.dbMaster.SelectMaster()
		assert.NoError(t, err, "Expected no error selecting master")
		require.NotNil(t, master, "Expected master to not be nil")
		assert.Equal(t, queuer2.worker.RID, master.WorkerRID, "Expected master RID to match worker 2 RID")
		assert.Equal(t, queuer2.worker.ID, master.WorkerID, "Expected master ID to match worker 2 ID")
		assert.Equal(t, queuer2.worker.Status, model.WorkerStatusRunning, "Expected worker status to be RUNNING")
		assert.Equal(t, *masterSettings2, master.Settings, "Expected master settings to match")

		// Check if the master updated_at is within the last 5 seconds,
		// as the ticker runs every 5 seconds and we waited 6 seconds.
		assert.GreaterOrEqual(t, master.UpdatedAt.Unix(), time.Now().Add(-5*time.Second).Unix(), "Expected master updated_at to be in the last 5 seconds")

		err = queuer2.Stop()
		assert.NoError(t, err, "Expected Stop to complete without error")
	})
}

func TestStartWithoutWorker(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	t.Run("Start Queuer without worker", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		startQueuer := func() {
			queuer.StartWithoutWorker(ctx, cancel, true)
		}
		assert.NotPanics(t, startQueuer, "Expected no panic when starting Queuer without worker")

		err := queuer.Stop()
		assert.NoError(t, err, "Expected Stop to complete without error")
	})

	t.Run("Start 2 Queuers with MasterSettings and cancel first", func(t *testing.T) {
		queuer1 := NewQueuer("test", 10)
		require.NotNil(t, queuer1, "Expected Queuer 1 to be created successfully")
		ctx1, cancel1 := context.WithCancel(context.Background())
		defer cancel1()
		masterSettings1 := &model.MasterSettings{
			MasterLockTimeout:  10 * time.Second, // Short timeout for test
			MasterPollInterval: 5 * time.Second,
			JobDeleteThreshold: 30 * 24 * time.Hour,
		}
		masterSettings1.SetDefault() // Set default values for comparison
		queuer1.StartWithoutWorker(ctx1, cancel1, true, masterSettings1)

		queuer2 := NewQueuer("test", 20)
		require.NotNil(t, queuer2, "Expected Queuer 2 to be created successfully")
		ctx2, cancel2 := context.WithCancel(context.Background())
		defer cancel2()
		masterSettings2 := &model.MasterSettings{
			MasterLockTimeout:  10 * time.Second, // Short timeout for test
			MasterPollInterval: 3 * time.Second,
			JobDeleteThreshold: 20 * 24 * time.Hour,
		}
		masterSettings2.SetDefault() // Set default values for comparison
		queuer2.StartWithoutWorker(ctx2, cancel2, true, masterSettings2)

		time.Sleep(6 * time.Second) // Wait for the ticker to run at least once

		master, err := queuer1.dbMaster.SelectMaster()
		assert.NoError(t, err, "Expected no error selecting master")
		require.NotNil(t, master, "Expected master to not be nil")
		assert.Equal(t, queuer1.worker.RID, master.WorkerRID, "Expected master RID to match worker RID")
		assert.Equal(t, queuer1.worker.ID, master.WorkerID, "Expected master ID to match worker ID")
		assert.Equal(t, *masterSettings1, master.Settings, "Expected master settings to match")

		// Check if the master updated_at is within the last 5 seconds,
		// as the ticker runs every 5 seconds and we waited 6 seconds.
		assert.GreaterOrEqual(t, master.UpdatedAt.Unix(), time.Now().Add(-5*time.Second).Unix(), "Expected master updated_at to be in the last 5 seconds")

		// Cancel the first queuer
		cancel1()

		time.Sleep(12 * time.Second) // Wait for master lock timeout (10s) + buffer

		master, err = queuer2.dbMaster.SelectMaster()
		assert.NoError(t, err, "Expected no error selecting master")
		require.NotNil(t, master, "Expected master to not be nil")
		assert.Equal(t, queuer2.worker.RID, master.WorkerRID, "Expected master RID to match worker 2 RID")
		assert.Equal(t, queuer2.worker.ID, master.WorkerID, "Expected master ID to match worker 2 ID")
		assert.Equal(t, *masterSettings2, master.Settings, "Expected master settings to match")

		// Check if the master updated_at is within the last 5 seconds,
		// as the ticker runs every 5 seconds and we waited 6 seconds.
		assert.GreaterOrEqual(t, master.UpdatedAt.Unix(), time.Now().Add(-5*time.Second).Unix(), "Expected master updated_at to be in the last 5 seconds")

		err = queuer2.Stop()
		assert.NoError(t, err, "Expected Stop to complete without error")
	})
}

func TestStop(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	queuer := NewQueuer("test", 10)
	require.NotNil(t, queuer, "Expected Queuer to be created successfully")

	err := queuer.Stop()
	assert.NoError(t, err, "Expected Stop to complete without error")
}

func TestHeartbeatTicker(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	t.Run("Heartbeat ticker starts successfully", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		defer queuer.Stop()

		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()

		err := queuer.heartbeatTicker(ctx)
		assert.NoError(t, err, "Expected heartbeat ticker to start without error")

		// Wait for heartbeat to run at least once (heartbeat interval is 30s, but we need to wait less)
		// Since we're using a short context timeout, the heartbeat should run immediately
		time.Sleep(100 * time.Millisecond)

		// The worker's updated_at should be refreshed (or at least attempted to be refreshed)
		// Since the heartbeat interval is 30s and we only wait 100ms, we mainly test that it starts
		assert.NotNil(t, queuer.worker, "Expected worker to still exist")
	})

	t.Run("Heartbeat ticker handles nil worker gracefully", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		defer queuer.Stop()

		// Temporarily set worker to nil to test graceful handling
		queuer.workerMu.Lock()
		originalWorker := queuer.worker
		queuer.worker = nil
		queuer.workerMu.Unlock()

		ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
		defer cancel()

		err := queuer.heartbeatTicker(ctx)
		assert.NoError(t, err, "Expected heartbeat ticker to start without error even with nil worker")

		// Restore worker
		queuer.workerMu.Lock()
		queuer.worker = originalWorker
		queuer.workerMu.Unlock()
	})

	t.Run("Heartbeat ticker stops when context is cancelled", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		defer queuer.Stop()

		ctx, cancel := context.WithCancel(context.Background())

		err := queuer.heartbeatTicker(ctx)
		assert.NoError(t, err, "Expected heartbeat ticker to start without error")

		// Cancel the context immediately
		cancel()

		// Give it a moment to stop
		time.Sleep(100 * time.Millisecond)

		// If we reach here without hanging, the ticker stopped properly
		assert.True(t, true, "Heartbeat ticker stopped when context was cancelled")
	})
}

// Long running task for testing stop functionality
func TaskLongRunning(duration int) error {
	time.Sleep(time.Duration(duration) * time.Second)
	return nil
}

func TestStopWorkerWithRunningJobs(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	t.Run("StopWorker cancels running jobs immediately", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		queuer.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer.WorkerPollInterval = 1 * time.Second

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel)

		// Add a long-running job
		job, err := queuer.AddJob(TaskLongRunning, nil, 10)
		require.NoError(t, err, "Expected AddJob to succeed")
		require.NotNil(t, job, "Expected job to be created")

		// Wait for job to start running
		time.Sleep(500 * time.Millisecond)

		// Verify job is running
		runningJob, err := queuer.GetJob(job.RID)
		require.NoError(t, err, "Expected GetJob to succeed")
		require.NotNil(t, runningJob, "Expected job to be retrieved")
		assert.Equal(t, model.JobStatusRunning, runningJob.Status, "Expected job status to be RUNNING")

		// Stop worker (non-gracefully) - should cancel running jobs
		err = queuer.StopWorker(queuer.worker.RID)
		assert.NoError(t, err, "Expected StopWorker to succeed")

		// Verify worker status is STOPPED
		worker, err := queuer.GetWorker(queuer.worker.RID)
		assert.NoError(t, err, "Expected GetWorker to succeed")
		require.NotNil(t, worker, "Expected worker to be retrieved")
		assert.Equal(t, "STOPPED", worker.Status, "Expected worker status to be STOPPED")

		// Wait for heartbeat ticker to detect STOPPED status and cancel jobs (1s heartbeat + buffer)
		time.Sleep(3 * time.Second)

		// Check if job was cancelled and moved to archive
		archivedJob, err := queuer.dbJob.SelectJobFromArchive(job.RID)
		if err != nil || archivedJob == nil {
			// Job might still be in main table if cancelled but not archived yet
			t.Skip("Job archival may be delayed, skipping archive check")
		}
		if archivedJob != nil {
			assert.Equal(t, model.JobStatusCancelled, archivedJob.Status, "Expected job status to be CANCELLED")
		}

		// Note: We don't call queuer.Stop() here because the heartbeat ticker
		// will automatically call Stop() when it detects the STOPPED status
	})

	t.Run("StopWorker with multiple running jobs", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		queuer.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer.WorkerPollInterval = 1 * time.Second

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel)

		// Add multiple long-running jobs
		job1, err := queuer.AddJob(TaskLongRunning, nil, 10)
		require.NoError(t, err, "Expected AddJob to succeed")
		job2, err := queuer.AddJob(TaskLongRunning, nil, 10)
		require.NoError(t, err, "Expected AddJob to succeed")
		job3, err := queuer.AddJob(TaskLongRunning, nil, 10)
		require.NoError(t, err, "Expected AddJob to succeed")

		// Wait for jobs to start running
		time.Sleep(500 * time.Millisecond)

		// Stop worker (non-gracefully)
		err = queuer.StopWorker(queuer.worker.RID)
		assert.NoError(t, err, "Expected StopWorker to succeed")

		// Wait for heartbeat ticker to detect STOPPED status and cancel jobs
		// Need extra time for jobs to be archived (heartbeat runs, jobs cancelled, then archived)
		time.Sleep(7 * time.Second)

		// Verify all jobs were cancelled
		cancelledCount := 0
		for _, jobRID := range []uuid.UUID{job1.RID, job2.RID, job3.RID} {
			archivedJob, err := queuer.dbJob.SelectJobFromArchive(jobRID)
			if err == nil && archivedJob != nil {
				if archivedJob.Status == model.JobStatusCancelled {
					cancelledCount++
				}
			}
		}
		// Jobs should be cancelled (archival might be delayed)
		if cancelledCount == 0 {
			t.Log("Warning: No jobs found in archive yet, they may still be processing")
		}

		// Note: We don't call queuer.Stop() here because the heartbeat ticker
		// will automatically call Stop() when it detects the STOPPED status
	})
}

func TestStopWorkerGracefullyWithRunningJobs(t *testing.T) {
	envs := map[string]string{
		"QUEUER_DB_HOST":     "localhost",
		"QUEUER_DB_PORT":     dbPort,
		"QUEUER_DB_DATABASE": "database",
		"QUEUER_DB_USERNAME": "user",
		"QUEUER_DB_PASSWORD": "password",
		"QUEUER_DB_SCHEMA":   "public",
		"QUEUER_DB_SSLMODE":  "disable",
	}
	for key, value := range envs {
		t.Setenv(key, value)
	}

	t.Run("StopWorkerGracefully waits for running jobs to finish", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		queuer.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer.WorkerPollInterval = 1 * time.Second

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel)

		// Add a job with moderate duration
		job, err := queuer.AddJob(TaskLongRunning, nil, 3)
		require.NoError(t, err, "Expected AddJob to succeed")
		require.NotNil(t, job, "Expected job to be created")

		// Wait for job to start running
		time.Sleep(500 * time.Millisecond)

		// Verify job is running
		runningJob, err := queuer.GetJob(job.RID)
		require.NoError(t, err, "Expected GetJob to succeed")
		require.NotNil(t, runningJob, "Expected job to be retrieved")
		assert.Equal(t, model.JobStatusRunning, runningJob.Status, "Expected job status to be RUNNING")

		// Stop worker gracefully - should allow running jobs to finish
		err = queuer.StopWorkerGracefully(queuer.worker.RID)
		assert.NoError(t, err, "Expected StopWorkerGracefully to succeed")

		// Verify worker status is STOPPING
		worker, err := queuer.GetWorker(queuer.worker.RID)
		assert.NoError(t, err, "Expected GetWorker to succeed")
		require.NotNil(t, worker, "Expected worker to be retrieved")
		assert.Equal(t, "STOPPING", worker.Status, "Expected worker status to be STOPPING")

		// Wait for job to finish naturally and then heartbeat to process graceful stop
		// Job takes 3s + heartbeat ticker checks every 1s
		time.Sleep(5 * time.Second)

		// Check if job finished successfully and moved to archive
		archivedJob, err := queuer.dbJob.SelectJobFromArchive(job.RID)
		if err != nil || archivedJob == nil {
			t.Skip("Job may not be archived yet or queuer stopped before archival")
		}
		if archivedJob != nil {
			// Job should have completed (either SUCCEEDED or FAILED due to early shutdown)
			assert.Contains(t, []string{model.JobStatusSucceeded, model.JobStatusFailed}, archivedJob.Status,
				"Expected job to complete (graceful stop should wait for completion)")
		}

		// Note: We don't call queuer.Stop() here because the heartbeat ticker
		// will automatically call Stop() when it detects all jobs are done
	})

	t.Run("StopWorkerGracefully does not accept new jobs", func(t *testing.T) {
		queuer := NewQueuer("test", 10)
		require.NotNil(t, queuer, "Expected Queuer to be created successfully")
		queuer.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer.WorkerPollInterval = 500 * time.Millisecond

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel)

		// Wait for initial heartbeat to complete
		time.Sleep(100 * time.Millisecond)

		// Get initial worker state
		initialWorker, err := queuer.GetWorker(queuer.worker.RID)
		require.NoError(t, err, "Expected GetWorker to succeed")
		require.NotNil(t, initialWorker, "Expected worker to be retrieved")
		t.Logf("Initial worker status: %s, maxConcurrency: %d", initialWorker.Status, initialWorker.MaxConcurrency)

		// Stop worker gracefully
		err = queuer.StopWorkerGracefully(queuer.worker.RID)
		assert.NoError(t, err, "Expected StopWorkerGracefully to succeed")

		// Verify status is STOPPING
		stoppingWorker, err := queuer.GetWorker(queuer.worker.RID)
		require.NoError(t, err, "Expected GetWorker to succeed")
		require.NotNil(t, stoppingWorker, "Expected worker to be retrieved")
		t.Logf("After StopWorkerGracefully - status: %s, maxConcurrency: %d", stoppingWorker.Status, stoppingWorker.MaxConcurrency)
		assert.Equal(t, "STOPPING", stoppingWorker.Status, "Expected worker status to be STOPPING")

		// Wait for heartbeat to run and update maxConcurrency
		// With 500ms interval, wait at least 1.5 seconds for 2-3 heartbeats
		time.Sleep(1500 * time.Millisecond)

		// Verify worker maxConcurrency is 0 (but only if queuer is still running)
		worker, err := queuer.GetWorker(queuer.worker.RID)
		if err == nil && worker != nil {
			t.Logf("After waiting - status: %s, maxConcurrency: %d", worker.Status, worker.MaxConcurrency)
			assert.Equal(t, 0, worker.MaxConcurrency, "Expected worker maxConcurrency to be 0 when STOPPING")

			// Add a new job - it should not be picked up by the stopping worker
			job, err := queuer.AddJob(TaskLongRunning, nil, 2)
			if err == nil {
				// Wait a moment for polling
				time.Sleep(1 * time.Second)

				// Job should still be queued, not running (if we can still access it)
				queuedJob, err := queuer.GetJob(job.RID)
				if err == nil && queuedJob != nil {
					assert.Equal(t, "QUEUED", queuedJob.Status, "Expected job to remain QUEUED when worker is stopping")
				}
			}
		}

		// Note: We don't call queuer.Stop() here because the heartbeat ticker
		// will automatically call Stop() when it detects all jobs are done
	})

	t.Run("Compare graceful vs non-graceful stop with running job", func(t *testing.T) {
		// Test 1: Non-graceful stop - job should be cancelled
		queuer1 := NewQueuer("test-nongraceful", 10)
		require.NotNil(t, queuer1, "Expected Queuer 1 to be created successfully")
		queuer1.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer1.WorkerPollInterval = 1 * time.Second

		ctx1, cancel1 := context.WithCancel(context.Background())
		defer cancel1()
		queuer1.Start(ctx1, cancel1)

		job1, err := queuer1.AddJob(TaskLongRunning, nil, 5)
		require.NoError(t, err, "Expected AddJob to succeed")

		// Wait for job to start
		time.Sleep(500 * time.Millisecond)

		// Non-graceful stop
		err = queuer1.StopWorker(queuer1.worker.RID)
		assert.NoError(t, err, "Expected StopWorker to succeed")

		// Wait for heartbeat to process STOPPED status
		time.Sleep(3 * time.Second)

		// Job should be cancelled (if archived)
		archivedJob1, err := queuer1.dbJob.SelectJobFromArchive(job1.RID)
		if err == nil && archivedJob1 != nil {
			assert.Equal(t, model.JobStatusCancelled, archivedJob1.Status, "Expected non-graceful stop to CANCEL the job")
		}

		// Test 2: Graceful stop - job should succeed
		queuer2 := NewQueuer("test-graceful", 10)
		require.NotNil(t, queuer2, "Expected Queuer 2 to be created successfully")
		queuer2.AddTask(TaskLongRunning)
		// Set fast heartbeat interval for testing
		queuer2.WorkerPollInterval = 1 * time.Second

		ctx2, cancel2 := context.WithCancel(context.Background())
		defer cancel2()
		queuer2.Start(ctx2, cancel2)

		job2, err := queuer2.AddJob(TaskLongRunning, nil, 3)
		require.NoError(t, err, "Expected AddJob to succeed")

		// Wait for job to start
		time.Sleep(500 * time.Millisecond)

		// Graceful stop
		err = queuer2.StopWorkerGracefully(queuer2.worker.RID)
		assert.NoError(t, err, "Expected StopWorkerGracefully to succeed")

		// Wait for job to complete + heartbeat to process graceful stop
		time.Sleep(5 * time.Second)

		// Job should complete (may succeed or fail due to shutdown timing)
		archivedJob2, err := queuer2.dbJob.SelectJobFromArchive(job2.RID)
		if err == nil && archivedJob2 != nil {
			// Graceful stop should allow completion (not cancel mid-execution)
			assert.Contains(t, []string{model.JobStatusSucceeded, model.JobStatusFailed}, archivedJob2.Status,
				"Expected graceful stop to allow job completion (not cancel)")
		}

		// Note: Both queuers will be stopped automatically by their heartbeat tickers
	})
}

func TestJobArchiveRetention(t *testing.T) {
	t.Run("TimescaleDB retention policy", func(t *testing.T) {
		teardown, timescalePort, err := helper.MustStartTimescaleContainer()
		require.NoError(t, err)
		defer teardown(context.Background())

		queuer := NewQueuerWithDB("test-retention", 10, "", &helper.DatabaseConfiguration{
			Host: "localhost", Port: timescalePort, Database: "database",
			Username: "user", Password: "password", Schema: "public",
			SSLMode: "disable", WithTableDrop: true, WithTimescale: true,
		})
		queuer.AddTaskWithName(func() (string, error) { return "success", nil }, "test-task")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel, &model.MasterSettings{JobDeleteThreshold: 1 * time.Hour})
		time.Sleep(6 * time.Second)

		master, err := queuer.dbMaster.SelectMaster()
		require.NoError(t, err)
		assert.Equal(t, 1*time.Hour, master.Settings.JobDeleteThreshold)

		queuer.Stop()
	})

	t.Run("PostgreSQL trigger-based retention", func(t *testing.T) {
		teardown, postgresPort, err := helper.MustStartPostgresContainer()
		require.NoError(t, err)
		defer teardown(context.Background())

		queuer := NewQueuerWithDB("test-retention", 10, "", &helper.DatabaseConfiguration{
			Host: "localhost", Port: postgresPort, Database: "database",
			Username: "user", Password: "password", Schema: "public",
			SSLMode: "disable", WithTableDrop: true, WithTimescale: false,
		})
		queuer.AddTaskWithName(func() (string, error) { return "success", nil }, "test-task")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		queuer.Start(ctx, cancel, &model.MasterSettings{JobDeleteThreshold: 1 * time.Hour})
		time.Sleep(6 * time.Second)

		// Insert old job
		oldJobRID := uuid.New()
		_, err = queuer.DB.Exec(`INSERT INTO job_archive (id, rid, worker_id, worker_rid, task_name, status, updated_at, created_at)
			VALUES (1, $1, 1, $2, 'test-task', 'SUCCEEDED', CURRENT_TIMESTAMP - INTERVAL '2 hours', CURRENT_TIMESTAMP - INTERVAL '2 hours')`,
			oldJobRID, queuer.worker.RID)
		require.NoError(t, err)

		// Add 15 jobs to trigger cleanup (trigger runs every 10th insert)
		for i := 0; i < 15; i++ {
			_, err := queuer.AddJob("test-task", nil)
			require.NoError(t, err)
		}

		// Wait for jobs to complete and trigger to execute
		time.Sleep(5 * time.Second)

		// Verify old job was deleted
		deletedJob, err := queuer.dbJob.SelectJobFromArchive(oldJobRID)
		assert.Error(t, err)
		assert.Nil(t, deletedJob)

		queuer.Stop()
	})
}