fix: set leader node only when cache is warmed

cmd/run.go

@@ -170,6 +170,16 @@ This enables a CRD-driven approach to automated image updates with Argo CD.
 				return err
 			}
 
+			// isLeader is an atomic bool that will be set to true when the controller
+			// becomes the leader.
+			var isLeader atomic.Bool
+
+			// Start a goroutine that sets isLeader to true when leadership is acquired.
+			go func() {
+				<-mgr.Elected()
+				isLeader.Store(true)
+			}()
+
 			// Add the CacheWarmer as a Runnable to the manager.
 			setupLogger.Info("Adding cache warmer to the manager.")
 			warmupState := &WarmupStatus{
@@ -231,11 +241,11 @@ This enables a CRD-driven approach to automated image updates with Argo CD.
 			}
 
 			if err := mgr.AddReadyzCheck("warmup-check", func(req *http.Request) error {
-				if !warmupState.isCacheWarmed.Load() {
-					// If the cache is not yet warmed, the check fails.
+				// If we are the leader, we are ready only if the cache is warmed.
+				if isLeader.Load() && !warmupState.isCacheWarmed.Load() {
 					return fmt.Errorf("cache is not yet warmed")
 				}
-				// Once warmed, the check passes.
+				// If we are not the leader, we are always ready.
 				return nil
 			}); err != nil {
 				setupLogger.Error(err, "unable to set up ready check")
@@ -420,6 +430,8 @@ func (ws *WebhookServerRunnable) Start(ctx context.Context) error {
 	return ws.webhookServer.Start(ctx)
 }
 
+// NeedLeaderElection tells the manager that this runnable should only be
+// run on the leader replica.
 func (ws *WebhookServerRunnable) NeedLeaderElection() bool {
 	return true
 }

cmd/run_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"strconv"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -340,3 +341,68 @@ func TestReadyzCheckWithWarmupStatus(t *testing.T) {
 		assert.Equal(t, "cache is not yet warmed", err.Error())
 	})
 }
+
+// Assisted-by: Gemini AI
+// TestLeadershipAwareReadyzCheck verifies the behavior of the leadership-aware readiness probe.
+func TestLeadershipAwareReadyzCheck(t *testing.T) {
+	testCases := []struct {
+		name        string
+		isLeader    bool
+		isWarmed    bool
+		expectError bool
+		errorMsg    string
+	}{
+		{
+			name:        "Leader and not warmed",
+			isLeader:    true,
+			isWarmed:    false,
+			expectError: true,
+			errorMsg:    "cache is not yet warmed",
+		},
+		{
+			name:        "Leader and warmed",
+			isLeader:    true,
+			isWarmed:    true,
+			expectError: false,
+		},
+		{
+			name:        "Not leader and not warmed",
+			isLeader:    false,
+			isWarmed:    false,
+			expectError: false,
+		},
+		{
+			name:        "Not leader and warmed",
+			isLeader:    false,
+			isWarmed:    true,
+			expectError: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			var isLeader atomic.Bool
+			isLeader.Store(tc.isLeader)
+
+			status := &WarmupStatus{Done: make(chan struct{})}
+			status.isCacheWarmed.Store(tc.isWarmed)
+
+			// This is the readiness check logic from cmd/run.go
+			readinessCheck := func(req *http.Request) error {
+				if isLeader.Load() && !status.isCacheWarmed.Load() {
+					return fmt.Errorf("cache is not yet warmed")
+				}
+				return nil
+			}
+
+			err := readinessCheck(nil)
+
+			if tc.expectError {
+				assert.Error(t, err)
+				assert.Equal(t, tc.errorMsg, err.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}