query-orchestration/internal/backgroundtask/test_helpers.go

package backgroundtask

import (
	"context"
	"errors"
	"log/slog"
	"sync/atomic"
	"time"
)

// TestConfig holds configuration for test work functions.
type TestConfig struct {
	// WorkDuration is how long each work execution should take
	WorkDuration time.Duration

	// ShouldError indicates if the work function should return an error
	ShouldError bool

	// ErrorMessage is the error message to return when ShouldError is true
	ErrorMessage string

	// RespectContext indicates if the work function should respect context cancellation
	RespectContext bool

	// Counter tracks the number of executions
	Counter *atomic.Int64

	// MaxExecutions limits the number of executions (0 = unlimited)
	MaxExecutions int64
}

// NewTestConfig creates a default test configuration.
func NewTestConfig() *TestConfig {
	return &TestConfig{
		WorkDuration:   1 * time.Millisecond,
		ShouldError:    false,
		ErrorMessage:   "test error",
		RespectContext: true,
		Counter:        &atomic.Int64{},
		MaxExecutions:  0,
	}
}

// NewTestConfigMap creates a config map with a default test configuration.
// This is a convenience function for testing.
func NewTestConfigMap() map[string]any {
	return map[string]any{
		"testConfig": NewTestConfig(),
	}
}

// GetTestConfig extracts a TestConfig from a config map.
// Returns the config and true if found, nil and false otherwise.
func GetTestConfig(config map[string]any) (*TestConfig, bool) {
	cfgVal, ok := config["testConfig"]
	if !ok {
		return nil, false
	}
	cfg, ok := cfgVal.(*TestConfig)
	return cfg, ok
}

// TestWorkFunc is a configurable work function for testing.
// It increments the counter and simulates work based on the TestConfig stored in the config map.
func TestWorkFunc(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	cfg, ok := GetTestConfig(config)
	if !ok {
		return errors.New("testConfig not found or invalid type")
	}

	// Check max executions
	if cfg.MaxExecutions > 0 {
		current := cfg.Counter.Load()
		if current >= cfg.MaxExecutions {
			logger.Debug("max executions reached", "count", current)
			return nil
		}
	}

	// Increment counter
	count := cfg.Counter.Add(1)
	logger.Debug("test work executing", "count", count)

	// Simulate work duration
	if cfg.WorkDuration > 0 {
		if cfg.RespectContext {
			// Use select to respect context cancellation
			select {
			case <-time.After(cfg.WorkDuration):
				// Work completed
			case <-ctx.Done():
				logger.Debug("test work canceled", "count", count)
				return ctx.Err()
			}
		} else {
			// Sleep regardless of context
			time.Sleep(cfg.WorkDuration)
		}
	}

	// Return error if configured
	if cfg.ShouldError {
		logger.Debug("test work returning error", "count", count, "error", cfg.ErrorMessage)
		return errors.New(cfg.ErrorMessage)
	}

	logger.Debug("test work completed", "count", count)
	return nil
}

// FastWorkFunc is a minimal work function that just increments a counter.
// Useful for testing rapid executions.
func FastWorkFunc(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	cfg, ok := GetTestConfig(config)
	if !ok {
		return errors.New("testConfig not found or invalid type")
	}

	count := cfg.Counter.Add(1)
	logger.Debug("fast work executed", "count", count)
	return nil
}

// SlowWorkFunc simulates a long-running task that respects context cancellation.
// It checks context every 10ms while simulating work.
func SlowWorkFunc(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	cfg, ok := GetTestConfig(config)
	if !ok {
		return errors.New("testConfig not found or invalid type")
	}

	count := cfg.Counter.Add(1)
	logger.Debug("slow work starting", "count", count, "duration", cfg.WorkDuration)

	// Check context periodically during long work
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	deadline := time.Now().Add(cfg.WorkDuration)
	for time.Now().Before(deadline) {
		select {
		case <-ticker.C:
			// Continue working
		case <-ctx.Done():
			logger.Debug("slow work canceled", "count", count)
			return ctx.Err()
		}
	}

	logger.Debug("slow work completed", "count", count)
	return nil
}

// ErrorWorkFunc always returns an error after incrementing the counter.
func ErrorWorkFunc(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	cfg, ok := GetTestConfig(config)
	if !ok {
		return errors.New("testConfig not found or invalid type")
	}

	count := cfg.Counter.Add(1)
	logger.Debug("error work executed", "count", count)

	if cfg.ErrorMessage != "" {
		return errors.New(cfg.ErrorMessage)
	}
	return errors.New("error work function failed")
}

// BlockingWorkFunc blocks until context is canceled.
// Useful for testing shutdown behavior.
func BlockingWorkFunc(ctx context.Context, logger *slog.Logger, config map[string]any) error {
	cfg, ok := GetTestConfig(config)
	if !ok {
		return errors.New("testConfig not found or invalid type")
	}

	count := cfg.Counter.Add(1)
	logger.Debug("blocking work started", "count", count)

	<-ctx.Done()

	logger.Debug("blocking work canceled", "count", count)
	return ctx.Err()
}

// WaitForWorkCycles waits for the specified number of work cycles to complete.
// Returns true if the target was reached within the timeout, false otherwise.
func WaitForWorkCycles(runner *Runner, target int64, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	for time.Now().Before(deadline) {
		stats := runner.Stats()
		if stats.WorkCycles() >= target {
			return true
		}
		<-ticker.C
	}
	return false
}

// WaitForCondition waits for a condition function to return true.
// Returns true if the condition was met within the timeout, false otherwise.
func WaitForCondition(condition func() bool, timeout time.Duration) bool {
	deadline := time.Now().Add(timeout)
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()

	for time.Now().Before(deadline) {
		if condition() {
			return true
		}
		<-ticker.C
	}
	return false
}