internal/workers/pool_test.go

package workers_test

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"

	"github.com/AmaniQuery/amaniquery/internal/workers"
	"go.uber.org/zap"
)

// TestDefaultConfig tests default configuration
func TestDefaultConfig(t *testing.T) {
	cfg := workers.DefaultConfig()

	if cfg.MaxWorkers <= 0 {
		t.Error("Expected positive MaxWorkers")
	}
	if cfg.QueueSize <= 0 {
		t.Error("Expected positive QueueSize")
	}
}

// TestNewPool tests pool creation
func TestNewPool(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	cfg := workers.Config{
		MaxWorkers: 4,
		QueueSize:  100,
		Logger:     logger,
	}

	pool := workers.NewPool(cfg)
	if pool == nil {
		t.Fatal("Expected non-nil pool")
	}

	// Clean shutdown
	err := pool.Shutdown(5 * time.Second)
	if err != nil {
		t.Fatalf("Shutdown failed: %v", err)
	}
}

// TestPool_Submit tests basic task submission
func TestPool_Submit(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	var executed int32
	task := workers.Task{
		ID: "test-task",
		Execute: func(ctx context.Context) error {
			atomic.AddInt32(&executed, 1)
			return nil
		},
	}

	err := pool.Submit(task)
	if err != nil {
		t.Fatalf("Submit failed: %v", err)
	}

	// Wait for task to execute
	time.Sleep(100 * time.Millisecond)

	if atomic.LoadInt32(&executed) != 1 {
		t.Error("Expected task to be executed")
	}
}

// TestPool_SubmitWait tests synchronous task submission
func TestPool_SubmitWait(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	ctx := context.Background()
	var executed bool

	err := pool.SubmitWait(ctx, func(ctx context.Context) error {
		executed = true
		return nil
	})

	if err != nil {
		t.Fatalf("SubmitWait failed: %v", err)
	}
	if !executed {
		t.Error("Expected task to be executed")
	}
}

// TestPool_SubmitWait_Error tests error handling in SubmitWait
func TestPool_SubmitWait_Error(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	ctx := context.Background()
	expectedErr := errors.New("task error")

	err := pool.SubmitWait(ctx, func(ctx context.Context) error {
		return expectedErr
	})

	if err == nil {
		t.Fatal("Expected error from task")
	}
	if err.Error() != expectedErr.Error() {
		t.Errorf("Expected error '%v', got '%v'", expectedErr, err)
	}
}

// TestPool_SubmitWait_ContextCancellation tests context cancellation
func TestPool_SubmitWait_ContextCancellation(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 1,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

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

	err := pool.SubmitWait(ctx, func(ctx context.Context) error {
		time.Sleep(200 * time.Millisecond)
		return nil
	})

	if err == nil {
		t.Fatal("Expected context cancellation error")
	}
}

// TestPool_SubmitBatch tests batch task submission
func TestPool_SubmitBatch(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 4,
		QueueSize:  20,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	var counter int32
	tasks := make([]workers.Task, 5)
	for i := range tasks {
		tasks[i] = workers.Task{
			ID: string(rune('0' + i)),
			Execute: func(ctx context.Context) error {
				atomic.AddInt32(&counter, 1)
				return nil
			},
		}
	}

	results := pool.SubmitBatch(tasks)

	// Collect results
	var errorCount int
	for err := range results {
		if err != nil {
			errorCount++
		}
	}

	if errorCount > 0 {
		t.Errorf("Expected no errors, got %d", errorCount)
	}

	if atomic.LoadInt32(&counter) != 5 {
		t.Errorf("Expected 5 tasks executed, got %d", counter)
	}
}

// TestPool_Metrics tests metrics tracking
func TestPool_Metrics(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	// Submit some tasks
	for i := 0; i < 5; i++ {
		pool.Submit(workers.Task{
			ID: string(rune('0' + i)),
			Execute: func(ctx context.Context) error {
				return nil
			},
		})
	}

	// Wait for tasks to complete
	time.Sleep(200 * time.Millisecond)

	metrics := pool.Metrics()

	if metrics.MaxWorkers != 2 {
		t.Errorf("Expected MaxWorkers 2, got %d", metrics.MaxWorkers)
	}
	if metrics.QueueCapacity != 10 {
		t.Errorf("Expected QueueCapacity 10, got %d", metrics.QueueCapacity)
	}
	if metrics.CompletedTasks < 5 {
		t.Errorf("Expected at least 5 completed tasks, got %d", metrics.CompletedTasks)
	}
}

// TestPool_OnError tests error callback
func TestPool_OnError(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	var capturedErr error
	expectedErr := errors.New("task error")

	task := workers.Task{
		ID: "error-task",
		Execute: func(ctx context.Context) error {
			return expectedErr
		},
		OnError: func(err error) {
			capturedErr = err
		},
	}

	pool.Submit(task)

	// Wait for task to execute
	time.Sleep(100 * time.Millisecond)

	if capturedErr == nil {
		t.Fatal("Expected error to be captured")
	}
	if capturedErr.Error() != expectedErr.Error() {
		t.Errorf("Expected error '%v', got '%v'", expectedErr, capturedErr)
	}
}

// TestPool_PanicRecovery tests panic recovery in tasks
func TestPool_PanicRecovery(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	var capturedErr error

	task := workers.Task{
		ID: "panic-task",
		Execute: func(ctx context.Context) error {
			panic("intentional panic")
		},
		OnError: func(err error) {
			capturedErr = err
		},
	}

	pool.Submit(task)

	// Wait for task to execute
	time.Sleep(100 * time.Millisecond)

	if capturedErr == nil {
		t.Fatal("Expected panic to be recovered and converted to error")
	}
}

// TestPool_Shutdown tests graceful shutdown
func TestPool_Shutdown(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})

	// Submit some work before shutdown
	var executed int32
	pool.Submit(workers.Task{
		ID: "pre-shutdown",
		Execute: func(ctx context.Context) error {
			atomic.AddInt32(&executed, 1)
			return nil
		},
	})

	// Wait briefly for task to execute
	time.Sleep(50 * time.Millisecond)

	err := pool.Shutdown(5 * time.Second)
	if err != nil {
		t.Fatalf("Shutdown failed: %v", err)
	}

	if atomic.LoadInt32(&executed) != 1 {
		t.Error("Expected task to execute before shutdown")
	}
}

// TestPool_Resize tests dynamic resizing
func TestPool_Resize(t *testing.T) {
	logger, _ := zap.NewDevelopment()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 2,
		QueueSize:  10,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	initialMetrics := pool.Metrics()
	if initialMetrics.MaxWorkers != 2 {
		t.Errorf("Expected initial MaxWorkers 2, got %d", initialMetrics.MaxWorkers)
	}

	// Resize to larger
	pool.Resize(4)

	metrics := pool.Metrics()
	if metrics.MaxWorkers != 4 {
		t.Errorf("Expected MaxWorkers 4 after resize, got %d", metrics.MaxWorkers)
	}
}

// TestPoolError tests error type
func TestPoolError(t *testing.T) {
	err := workers.ErrPoolShutdown
	if err.Error() != "worker pool is shutdown" {
		t.Errorf("Expected 'worker pool is shutdown', got '%s'", err.Error())
	}

	err = workers.ErrQueueFull
	if err.Error() != "task queue is full" {
		t.Errorf("Expected 'task queue is full', got '%s'", err.Error())
	}
}

// BenchmarkPool_Submit benchmarks task submission
func BenchmarkPool_Submit(b *testing.B) {
	logger, _ := zap.NewProduction()
	pool := workers.NewPool(workers.Config{
		MaxWorkers: 8,
		QueueSize:  10000,
		Logger:     logger,
	})
	defer pool.Shutdown(5 * time.Second)

	task := workers.Task{
		Execute: func(ctx context.Context) error {
			return nil
		},
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		pool.Submit(task)
	}
}