Fix UT to verify runnable ordering.

Author: godwinpang

pkg/manager/manager_test.go

@@ -1931,8 +1931,10 @@ var _ = Describe("manger.Manager", func() {
 
 	It("should run warmup runnables before leader election is won", func() {
 		By("Creating channels to track execution order")
-		warmupCalled := make(chan struct{})
-		leaderElectionRunnableCalled := make(chan struct{})
+		warmupCalled := atomic.Bool{}
+		leaderElectionRunnableStarted := atomic.Bool{}
+		// warmupRunnable's WarmupFunc will block until this channel is closed
+		warmupRunnableWarmupBlockingChan := make(chan struct{})
 
 		By("Creating a manager with leader election enabled")
 		m, err := New(cfg, Options{
@@ -1946,31 +1948,29 @@ var _ = Describe("manger.Manager", func() {
 		})
 		Expect(err).NotTo(HaveOccurred())
 
-		// Override onStoppedLeading to prevent os.Exit
-		cm := m.(*controllerManager)
-		cm.onStoppedLeading = func() {}
-
 		By("Creating a runnable that implements WarmupRunnable interface")
 		// Create a warmup runnable
 		warmupRunnable := warmupRunnableFunc{
 			StartFunc: func(ctx context.Context) error {
 				// This is the main runnable that will be executed after leader election
+				// Block forever
 				<-ctx.Done()
 				return nil
 			},
 			WarmupFunc: func(ctx context.Context) error {
 				// This should be called during startup before leader election
-				close(warmupCalled)
+				warmupCalled.Store(true)
+				<-warmupRunnableWarmupBlockingChan
 				return nil
 			},
+			didWarmupFinish: make(chan bool),
 		}
 		Expect(m.Add(warmupRunnable)).To(Succeed())
 
 		By("Creating a runnable that requires leader election")
 		leaderElectionRunnable := leaderElectionRunnableFunc{
 			StartFunc: func(ctx context.Context) error {
-				// This will only be called after leader election is won
-				close(leaderElectionRunnableCalled)
+				leaderElectionRunnableStarted.Store(true)
 				<-ctx.Done()
 				return nil
 			},
@@ -1985,17 +1985,21 @@ var _ = Describe("manger.Manager", func() {
 
 		go func() {
 			defer GinkgoRecover()
-			Expect(m.Start(ctx)).NotTo(HaveOccurred())
+			Expect(m.Start(ctx)).To(Succeed())
 		}()
 
 		By("Waiting for the warmup runnable to be called")
-		<-warmupCalled
+		Eventually(warmupCalled.Load).Should(BeTrue())
+		Expect(leaderElectionRunnableStarted.Load()).To(BeFalse())
+
+		By("Closing the channel to unblock the warmup runnable")
+		close(warmupRunnableWarmupBlockingChan)
 
 		By("Waiting for leader election to be won")
 		<-m.Elected()
 
 		By("Verifying the leader election runnable is called after election")
-		<-leaderElectionRunnableCalled
+		Eventually(leaderElectionRunnableStarted.Load).Should(BeTrue())
 	})
 })
 

pkg/manager/runnable_group_test.go

@@ -377,39 +377,48 @@ var _ warmupRunnable = &warmupRunnableFunc{}
 // warmupRunnableFunc is a helper struct that implements WarmupRunnable
 // for testing purposes.
 type warmupRunnableFunc struct {
-	StartFunc  func(context.Context) error
-	WarmupFunc func(context.Context) error
+	StartFunc       func(context.Context) error
+	WarmupFunc      func(context.Context) error
+	didWarmupFinish chan bool
 }
 
 func (r warmupRunnableFunc) Start(ctx context.Context) error {
 	return r.StartFunc(ctx)
 }
 
 func (r warmupRunnableFunc) Warmup(ctx context.Context) error {
-	return r.WarmupFunc(ctx)
+	err := r.WarmupFunc(ctx)
+	r.didWarmupFinish <- (err == nil)
+	return err
 }
 
 func (r warmupRunnableFunc) WaitForWarmupComplete(ctx context.Context) bool {
-	return true
+	return <-r.didWarmupFinish
 }
 
+var _ LeaderElectionRunnable = &leaderElectionAndWarmupRunnable{}
+var _ warmupRunnable = &leaderElectionAndWarmupRunnable{}
+
 // leaderElectionAndWarmupRunnable implements both WarmupRunnable and LeaderElectionRunnable
 type leaderElectionAndWarmupRunnable struct {
 	StartFunc              func(context.Context) error
 	WarmupFunc             func(context.Context) error
 	NeedLeaderElectionFunc func() bool
+	didWarmupFinish        chan bool
 }
 
 func (r leaderElectionAndWarmupRunnable) Start(ctx context.Context) error {
 	return r.StartFunc(ctx)
 }
 
 func (r leaderElectionAndWarmupRunnable) Warmup(ctx context.Context) error {
-	return r.WarmupFunc(ctx)
+	err := r.WarmupFunc(ctx)
+	r.didWarmupFinish <- (err == nil)
+	return err
 }
 
 func (r leaderElectionAndWarmupRunnable) WaitForWarmupComplete(ctx context.Context) bool {
-	return true
+	return <-r.didWarmupFinish
 }
 
 func (r leaderElectionAndWarmupRunnable) NeedLeaderElection() bool {