From 10cd54bee4657d3b9f1f6a7a3d770b7ad602ab3b Mon Sep 17 00:00:00 2001
From: Dan Winship <danwinship@redhat.com>
Date: Wed, 19 Mar 2025 16:52:03 -0400
Subject: [PATCH] Streamling the node/zone-picking logic

---
 test/e2e/network/traffic_distribution.go | 28 +++++++++---------------
 1 file changed, 10 insertions(+), 18 deletions(-)

diff --git a/test/e2e/network/traffic_distribution.go b/test/e2e/network/traffic_distribution.go
index de38d477cf2..34c08ba0e47 100644
--- a/test/e2e/network/traffic_distribution.go
+++ b/test/e2e/network/traffic_distribution.go
@@ -110,31 +110,23 @@ var _ = common.SIGDescribe("Traffic Distribution", func() {
 	////////////////////////////////////////////////////////////////////////////
 
 	ginkgo.It("should route traffic to an endpoint in the same zone when using PreferClose", func(ctx context.Context) {
-
 		ginkgo.By("finding 3 zones with schedulable nodes")
-		allZonesSet, err := e2enode.GetSchedulableClusterZones(ctx, c)
-		framework.ExpectNoError(err)
-		if len(allZonesSet) < 3 {
-			framework.Failf("got %d zones with schedulable nodes, want atleast 3 zones with schedulable nodes", len(allZonesSet))
-		}
-		zones := allZonesSet.UnsortedList()[:3]
-
-		ginkgo.By(fmt.Sprintf("finding a node in each of the chosen 3 zones %v", zones))
 		nodeList, err := e2enode.GetReadySchedulableNodes(ctx, c)
 		framework.ExpectNoError(err)
 		nodeForZone := make(map[string]*v1.Node)
-		for _, zone := range zones {
-			found := false
-			for _, node := range nodeList.Items {
-				if zone == node.Labels[v1.LabelTopologyZone] {
-					found = true
-					nodeForZone[zone] = &node
-				}
+		for _, node := range nodeList.Items {
+			zone := node.Labels[v1.LabelTopologyZone]
+			if nodeForZone[zone] != nil {
+				continue
 			}
-			if !found {
-				framework.Failf("could not find a node in zone %q; nodes=\n%v", zone, format.Object(nodeList, 1 /* indent one level */))
+			nodeForZone[zone] = &node
+			if len(nodeForZone) == 3 {
+				break
 			}
 		}
+		if len(nodeForZone) < 3 {
+			e2eskipper.Skipf("got %d zones with schedulable nodes, need at least 3", len(nodeForZone))
+		}
 
 		var clientPods []*clientPod
 		var serverPods []*serverPod