Don't assume there is only one replica in e2e tests

For the upgrade e2e tests, don't assume there is onle one replica.
Get the number of replicas from the deployment and wait for the
deployment to have that many available. Use the lease to determine
the leader pod and reference that.

Note that the name format of leases for operator-controller and catalogd
are quite different; this doesn't change that, as it may have an impact
on the upgrade test itself.

Signed-off-by: Todd Short <tshort@redhat.com>
Assisted-by: Claude Code

Author: tmshort

test/upgrade-e2e/post_upgrade_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	appsv1 "k8s.io/api/apps/v1"
+	coordinationv1 "k8s.io/api/coordination/v1"
 	corev1 "k8s.io/api/core/v1"
 	apimeta "k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -195,13 +196,14 @@ func TestClusterExtensionAfterOLMUpgrade(t *testing.T) {
 
 // waitForDeployment checks that the updated deployment with the given app.kubernetes.io/name label
 // has reached the desired number of replicas and that the number pods matches that number
-// i.e. no old pods remain. It will return a pointer to the first pod. This is only necessary
+// i.e. no old pods remain. It will return a pointer to the leader pod. This is only necessary
 // to facilitate the mitigation put in place for https://github.com/operator-framework/operator-controller/issues/1626
 func waitForDeployment(t *testing.T, ctx context.Context, controlPlaneLabel string) *corev1.Pod {
 	deploymentLabelSelector := labels.Set{"app.kubernetes.io/name": controlPlaneLabel}.AsSelector()
 
 	t.Log("Checking that the deployment is updated")
 	var desiredNumReplicas int32
+	var deploymentNamespace string
 	require.EventuallyWithT(t, func(ct *assert.CollectT) {
 		var managerDeployments appsv1.DeploymentList
 		require.NoError(ct, c.List(ctx, &managerDeployments, client.MatchingLabelsSelector{Selector: deploymentLabelSelector}))
@@ -215,15 +217,51 @@ func waitForDeployment(t *testing.T, ctx context.Context, controlPlaneLabel stri
 				managerDeployment.Status.ReadyReplicas == *managerDeployment.Spec.Replicas,
 		)
 		desiredNumReplicas = *managerDeployment.Spec.Replicas
+		deploymentNamespace = managerDeployment.Namespace
 	}, time.Minute, time.Second)
 
 	var managerPods corev1.PodList
 	t.Logf("Ensure the number of remaining pods equal the desired number of replicas (%d)", desiredNumReplicas)
 	require.EventuallyWithT(t, func(ct *assert.CollectT) {
 		require.NoError(ct, c.List(ctx, &managerPods, client.MatchingLabelsSelector{Selector: deploymentLabelSelector}))
-		require.Len(ct, managerPods.Items, 1)
+		require.Len(ct, managerPods.Items, int(desiredNumReplicas))
 	}, time.Minute, time.Second)
-	return &managerPods.Items[0]
+
+	// Find the leader pod by checking the lease
+	t.Log("Finding the leader pod")
+	// Map component labels to their leader election lease names
+	leaseNames := map[string]string{
+		"catalogd":            "catalogd-operator-lock",
+		"operator-controller": "9c4404e7.operatorframework.io",
+	}
+
+	leaseName, ok := leaseNames[controlPlaneLabel]
+	if !ok {
+		t.Fatalf("Unknown control plane component: %s", controlPlaneLabel)
+	}
+
+	var leaderPod *corev1.Pod
+	require.EventuallyWithT(t, func(ct *assert.CollectT) {
+		var lease coordinationv1.Lease
+		require.NoError(ct, c.Get(ctx, types.NamespacedName{Name: leaseName, Namespace: deploymentNamespace}, &lease))
+		require.NotNil(ct, lease.Spec.HolderIdentity)
+
+		leaderIdentity := *lease.Spec.HolderIdentity
+		// The lease holder identity format is: <pod-name>_<leader-election-id-suffix>
+		// Extract just the pod name by splitting on '_'
+		podName := strings.Split(leaderIdentity, "_")[0]
+
+		// Find the pod with matching name
+		for i := range managerPods.Items {
+			if managerPods.Items[i].Name == podName {
+				leaderPod = &managerPods.Items[i]
+				break
+			}
+		}
+		require.NotNil(ct, leaderPod, "leader pod not found with identity: %s (pod name: %s)", leaderIdentity, podName)
+	}, time.Minute, time.Second)
+
+	return leaderPod
 }
 
 func watchPodLogsForSubstring(ctx context.Context, pod *corev1.Pod, substrings ...string) (bool, error) {