fetch_ml/internal/queue/queue_test.go

package queue

import (
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestTaskQueue(t *testing.T) {
	// Start miniredis
	s, err := miniredis.Run()
	if err != nil {
		t.Fatalf("failed to start miniredis: %v", err)
	}
	defer s.Close()

	// Create TaskQueue
	cfg := Config{
		RedisAddr:            s.Addr(),
		MetricsFlushInterval: 10 * time.Millisecond, // Fast flush for testing
	}
	tq, err := NewTaskQueue(cfg)
	assert.NoError(t, err)
	defer tq.Close()

	t.Run("AddTask", func(t *testing.T) {
		task := &Task{
			ID:        "task-1",
			JobName:   "job-1",
			Status:    "queued",
			Priority:  10,
			CreatedAt: time.Now(),
		}
		err = tq.AddTask(task)
		assert.NoError(t, err)

		// Verify task is in Redis
		// Check ZSET
		score, err := s.ZScore(TaskQueueKey, "task-1")
		assert.NoError(t, err)
		assert.Equal(t, float64(10), score)
	})

	t.Run("GetNextTask", func(t *testing.T) {
		// Add another task
		task := &Task{
			ID:        "task-2",
			JobName:   "job-2",
			Status:    "queued",
			Priority:  20, // Higher priority
			CreatedAt: time.Now(),
		}
		err = tq.AddTask(task)
		assert.NoError(t, err)

		// Should get task-2 first due to higher priority
		nextTask, err := tq.GetNextTask()
		assert.NoError(t, err)
		assert.NotNil(t, nextTask)
		assert.Equal(t, "task-2", nextTask.ID)

		// Verify task is removed from ZSET
		_, err = tq.client.ZScore(tq.ctx, TaskQueueKey, "task-2").Result()
		assert.Equal(t, redis.Nil, err)
	})

	t.Run("GetNextTaskWithLease", func(t *testing.T) {
		task := &Task{
			ID:        "task-lease",
			JobName:   "job-lease",
			Status:    "queued",
			Priority:  15,
			CreatedAt: time.Now(),
		}
		err := tq.AddTask(task)
		require.NoError(t, err)

		workerID := "worker-1"
		leaseDuration := 1 * time.Minute

		leasedTask, err := tq.GetNextTaskWithLease(workerID, leaseDuration)
		require.NoError(t, err)
		require.NotNil(t, leasedTask)
		assert.Equal(t, "task-lease", leasedTask.ID)
		assert.Equal(t, workerID, leasedTask.LeasedBy)
		assert.NotNil(t, leasedTask.LeaseExpiry)
		assert.True(t, leasedTask.LeaseExpiry.After(time.Now()))
	})

	t.Run("RenewLease", func(t *testing.T) {
		taskID := "task-lease"
		workerID := "worker-1"

		// Get initial expiry
		task, err := tq.GetTask(taskID)
		require.NoError(t, err)
		initialExpiry := task.LeaseExpiry

		// Wait a bit
		time.Sleep(10 * time.Millisecond)

		// Renew lease
		err = tq.RenewLease(taskID, workerID, 1*time.Minute)
		require.NoError(t, err)

		// Verify expiry updated
		task, err = tq.GetTask(taskID)
		require.NoError(t, err)
		assert.True(t, task.LeaseExpiry.After(*initialExpiry))
	})

	t.Run("ReleaseLease", func(t *testing.T) {
		taskID := "task-lease"
		workerID := "worker-1"

		err := tq.ReleaseLease(taskID, workerID)
		require.NoError(t, err)

		task, err := tq.GetTask(taskID)
		require.NoError(t, err)
		assert.Nil(t, task.LeaseExpiry)
		assert.Empty(t, task.LeasedBy)
	})

	t.Run("RetryTask", func(t *testing.T) {
		task := &Task{
			ID:         "task-retry",
			JobName:    "job-retry",
			Status:     "failed",
			Priority:   10,
			CreatedAt:  time.Now(),
			MaxRetries: 3,
			RetryCount: 0,
			Error:      "some transient error",
		}

		// Add task directly to verify retry logic
		err := tq.AddTask(task)
		require.NoError(t, err)

		// Simulate failure and retry
		task.Error = "connection timeout"
		err = tq.RetryTask(task)
		require.NoError(t, err)

		// Verify task updated
		updatedTask, err := tq.GetTask(task.ID)
		require.NoError(t, err)
		assert.Equal(t, 1, updatedTask.RetryCount)
		assert.Equal(t, "queued", updatedTask.Status)
		assert.Empty(t, updatedTask.Error)
		assert.Equal(t, "connection timeout", updatedTask.LastError)
		assert.NotNil(t, updatedTask.NextRetry)
	})

	t.Run("DLQ", func(t *testing.T) {
		task := &Task{
			ID:         "task-dlq",
			JobName:    "job-dlq",
			Status:     "failed",
			Priority:   10,
			CreatedAt:  time.Now(),
			MaxRetries: 1,
			RetryCount: 1, // Already at max retries
			Error:      "fatal error",
		}

		err := tq.AddTask(task)
		require.NoError(t, err)

		// Retry should move to DLQ
		err = tq.RetryTask(task)
		require.NoError(t, err)

		// Verify removed from main queue
		_, err = tq.client.ZScore(tq.ctx, TaskQueueKey, task.ID).Result()
		assert.Equal(t, redis.Nil, err)

		// Verify in DLQ
		dlqKey := "task:dlq:" + task.ID
		exists := s.Exists(dlqKey)
		assert.True(t, exists)

		// Verify DLQ content
		val, err := s.Get(dlqKey)
		require.NoError(t, err)
		assert.Contains(t, val, "max retries exceeded")
	})
}