Reloader/test/e2e/utils/workload_adapter.go

package utils

import (
	"context"
	"time"

	"k8s.io/client-go/kubernetes"
)

// WorkloadType represents the type of Kubernetes workload.
type WorkloadType string

const (
	WorkloadDeployment       WorkloadType = "Deployment"
	WorkloadDaemonSet        WorkloadType = "DaemonSet"
	WorkloadStatefulSet      WorkloadType = "StatefulSet"
	WorkloadCronJob          WorkloadType = "CronJob"
	WorkloadJob              WorkloadType = "Job"
	WorkloadArgoRollout      WorkloadType = "ArgoRollout"
	WorkloadDeploymentConfig WorkloadType = "DeploymentConfig"
)

// ReloadStrategy represents the reload strategy used by Reloader.
type ReloadStrategy string

const (
	StrategyAnnotations ReloadStrategy = "annotations"
	StrategyEnvVars     ReloadStrategy = "envvars"
)

// WorkloadConfig holds configuration for workload creation.
type WorkloadConfig struct {
	ConfigMapName          string
	SecretName             string
	SPCName                string
	Annotations            map[string]string // Annotations for workload metadata (e.g., Deployment.metadata.annotations)
	PodTemplateAnnotations map[string]string // Annotations for pod template metadata (e.g., Deployment.spec.template.metadata.annotations)
	UseConfigMapEnvFrom    bool
	UseSecretEnvFrom       bool
	UseConfigMapVolume     bool
	UseSecretVolume        bool
	UseProjectedVolume     bool
	UseConfigMapKeyRef     bool
	UseSecretKeyRef        bool
	UseInitContainer       bool
	UseInitContainerVolume bool
	UseCSIVolume           bool
	UseInitContainerCSI    bool
	ConfigMapKey           string
	SecretKey              string
	EnvVarName             string
	MultipleContainers     int
}

// WorkloadAdapter provides a unified interface for all workload types.
// This allows tests to be parameterized across different workload types.
type WorkloadAdapter interface {
	// Type returns the workload type.
	Type() WorkloadType

	// Create creates the workload with the given config.
	Create(ctx context.Context, namespace, name string, cfg WorkloadConfig) error

	// Delete removes the workload.
	Delete(ctx context.Context, namespace, name string) error

	// WaitReady waits for the workload to be ready.
	WaitReady(ctx context.Context, namespace, name string, timeout time.Duration) error

	// WaitReloaded waits for the workload to have the reload annotation.
	// Returns true if the annotation was found, false if timeout occurred.
	// It captures the baseline annotation value at call time, which races with Reloader
	// if the reload trigger happened earlier — prefer WaitReloadedFrom in that case.
	WaitReloaded(ctx context.Context, namespace, name, annotationKey string, timeout time.Duration) (bool, error)

	// WaitReloadedFrom waits for the reload annotation to be present with a value different
	// from priorValue. Capture priorValue (via GetPodTemplateAnnotation) BEFORE performing
	// the change that triggers the reload; capturing it afterwards can observe the already
	// reloaded value and then wait for a further change that never comes.
	WaitReloadedFrom(ctx context.Context, namespace, name, annotationKey, priorValue string, timeout time.Duration) (bool, error)

	// WaitEnvVar waits for the workload to have a STAKATER_ env var (for envvars strategy).
	// Returns true if the env var was found, false if timeout occurred.
	// It captures the baseline env var value at call time, which races with Reloader
	// if the reload trigger happened earlier — prefer WaitEnvVarFrom in that case.
	WaitEnvVar(ctx context.Context, namespace, name, prefix string, timeout time.Duration) (bool, error)

	// WaitEnvVarFrom waits for a STAKATER_ env var whose value differs from priorValue.
	// Capture priorValue BEFORE performing the change that triggers the reload
	// (an empty priorValue means the env var is expected to appear).
	WaitEnvVarFrom(ctx context.Context, namespace, name, prefix, priorValue string, timeout time.Duration) (bool, error)

	// SupportsEnvVarStrategy returns true if the workload supports env var reload strategy.
	// CronJob does not support this as it uses job creation instead.
	SupportsEnvVarStrategy() bool

	// RequiresSpecialHandling returns true for workloads that need special handling.
	// For example, CronJob triggers a new job instead of rolling restart.
	RequiresSpecialHandling() bool

	// GetPodTemplateAnnotation returns the value of a pod template annotation.
	// This is useful for tests that need to compare annotation values before/after updates.
	GetPodTemplateAnnotation(ctx context.Context, namespace, name, annotationKey string) (string, error)
}

// Pausable is implemented by workloads that support pause/unpause.
// Currently only Deployment supports this capability.
type Pausable interface {
	WaitPaused(ctx context.Context, namespace, name, annotationKey string, timeout time.Duration) (bool, error)
	WaitUnpaused(ctx context.Context, namespace, name, annotationKey string, timeout time.Duration) (bool, error)
}

// Recreatable is implemented by workloads that are recreated instead of updated.
// Currently only Job supports this capability (Jobs are immutable, so Reloader recreates them).
type Recreatable interface {
	GetOriginalUID(ctx context.Context, namespace, name string) (string, error)
	WaitRecreated(ctx context.Context, namespace, name, originalUID string, timeout time.Duration) (string, bool, error)
}

// JobTriggerer is implemented by workloads that trigger jobs on reload.
// Currently only CronJob supports this capability.
type JobTriggerer interface {
	WaitForTriggeredJob(ctx context.Context, namespace, name string, timeout time.Duration) (bool, error)
}

// RestartAtSupporter is implemented by workloads that support the restartAt field.
// Currently only ArgoRollout supports this capability.
type RestartAtSupporter interface {
	WaitRestartAt(ctx context.Context, namespace, name string, timeout time.Duration) (bool, error)
}

// AdapterRegistry holds adapters for all workload types.
type AdapterRegistry struct {
	kubeClient kubernetes.Interface
	adapters   map[WorkloadType]WorkloadAdapter
}

// NewAdapterRegistry creates a new adapter registry with all standard adapters.
func NewAdapterRegistry(kubeClient kubernetes.Interface) *AdapterRegistry {
	r := &AdapterRegistry{
		kubeClient: kubeClient,
		adapters:   make(map[WorkloadType]WorkloadAdapter),
	}

	r.adapters[WorkloadDeployment] = NewDeploymentAdapter(kubeClient)
	r.adapters[WorkloadDaemonSet] = NewDaemonSetAdapter(kubeClient)
	r.adapters[WorkloadStatefulSet] = NewStatefulSetAdapter(kubeClient)
	r.adapters[WorkloadCronJob] = NewCronJobAdapter(kubeClient)
	r.adapters[WorkloadJob] = NewJobAdapter(kubeClient)

	return r
}

// RegisterAdapter registers a custom adapter for a workload type.
func (r *AdapterRegistry) RegisterAdapter(adapter WorkloadAdapter) {
	r.adapters[adapter.Type()] = adapter
}

// Get returns the adapter for the given workload type.
// Returns nil if the adapter is not registered.
func (r *AdapterRegistry) Get(wt WorkloadType) WorkloadAdapter {
	return r.adapters[wt]
}

// GetStandardWorkloads returns the standard workload types that are always available.
func (r *AdapterRegistry) GetStandardWorkloads() []WorkloadType {
	return []WorkloadType{
		WorkloadDeployment,
		WorkloadDaemonSet,
		WorkloadStatefulSet,
	}
}

// GetAllWorkloads returns all registered workload types in a canonical, deterministic order.
// Map iteration order in Go is non-deterministic, so this uses a fixed ordering to ensure
// consistent test parameterization across runs.
func (r *AdapterRegistry) GetAllWorkloads() []WorkloadType {
	canonical := []WorkloadType{
		WorkloadDeployment, WorkloadDaemonSet, WorkloadStatefulSet,
		WorkloadCronJob, WorkloadJob, WorkloadArgoRollout, WorkloadDeploymentConfig,
	}
	result := make([]WorkloadType, 0, len(r.adapters))
	for _, wt := range canonical {
		if _, ok := r.adapters[wt]; ok {
			result = append(result, wt)
		}
	}
	return result
}

// GetEnvVarWorkloads returns workload types that support env var reload strategy.
func (r *AdapterRegistry) GetEnvVarWorkloads() []WorkloadType {
	result := make([]WorkloadType, 0)
	for wt, adapter := range r.adapters {
		if adapter.SupportsEnvVarStrategy() {
			result = append(result, wt)
		}
	}
	return result
}