From 420ad059f662c3414bee148a5151f5f075c62af1 Mon Sep 17 00:00:00 2001
From: Izabela Gomes <igomes@vmware.com>
Date: Wed, 5 Jun 2024 09:09:55 -0400
Subject: [PATCH] use helper WaitForDeploymentScale

---
 test/ha/workload_test.go | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/test/ha/workload_test.go b/test/ha/workload_test.go
index 5e03ebb11fb2..9f9441335e40 100644
--- a/test/ha/workload_test.go
+++ b/test/ha/workload_test.go
@@ -29,6 +29,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/watch"
+	pkgTest "knative.dev/pkg/test"
 	"knative.dev/serving/pkg/apis/autoscaling"
 	"knative.dev/serving/pkg/apis/serving"
 	rtesting "knative.dev/serving/pkg/testing/v1"
@@ -40,6 +41,7 @@ import (
 const (
 	minimumNumberOfReplicas = 2
 	maximumNumberOfReplicas = 2
+	deploymentSuffix        = "-deployment"
 )
 
 func TestActivatorNotInRequestPath(t *testing.T) {
@@ -100,8 +102,12 @@ func testUptimeDuringUserPodDeletion(t *testing.T, ctx context.Context, clients
 	prober.Spawn(resources.Service.Status.URL.URL())
 	defer assertSLO(t, prober, 1)
 
-	// Get user pods
+	deploymentName := names.Revision + deploymentSuffix
+	if err := pkgTest.WaitForDeploymentScale(ctx, clients.KubeClient, deploymentName, test.ServingFlags.TestNamespace, minimumNumberOfReplicas); err != nil {
+		t.Fatalf("Deployment %s not scaled to %d: %v", deploymentName, minimumNumberOfReplicas, err)
+	}
 
+	// Get user pods
 	selector := labels.SelectorFromSet(labels.Set{
 		serving.ServiceLabelKey: names.Service,
 	})
@@ -110,10 +116,6 @@ func testUptimeDuringUserPodDeletion(t *testing.T, ctx context.Context, clients
 		t.Fatalf("Unable to get pods: %v", err)
 	}
 
-	if !(len(pods.Items) == minimumNumberOfReplicas) {
-		t.Fatalf("Expected to have %d user pod(s) running, but found %d.", minimumNumberOfReplicas, len(pods.Items))
-	}
-
 	t.Logf("Watching user pods")
 	var wg sync.WaitGroup
 	wg.Add(2)
@@ -137,8 +139,8 @@ func testUptimeDuringUserPodDeletion(t *testing.T, ctx context.Context, clients
 		t.Fatalf("Unable to get pods: %v", err)
 	}
 
-	if !(len(newPods.Items) == minimumNumberOfReplicas) {
-		t.Errorf("Expected to have %d user pod(s) running, but found %d.", minimumNumberOfReplicas, len(newPods.Items))
+	if err := pkgTest.WaitForDeploymentScale(ctx, clients.KubeClient, deploymentName, test.ServingFlags.TestNamespace, minimumNumberOfReplicas); err != nil {
+		t.Errorf("Deployment %s not scaled to %d: %v", deploymentName, minimumNumberOfReplicas, err)
 	}
 }