From 0301fc1e0ea73f55f264ef3a9af727a34c8ab581 Mon Sep 17 00:00:00 2001
From: Satish Matti <smatti@google.com>
Date: Thu, 21 Nov 2019 21:14:23 -0800
Subject: [PATCH] Update finalizer tests to validate against an ingress
 finalizer

---
 cmd/e2e-test/finalizer_test.go    | 20 ++++++++------------
 cmd/e2e-test/upgrade/finalizer.go |  6 +++++-
 pkg/e2e/helpers.go                | 28 +++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/cmd/e2e-test/finalizer_test.go b/cmd/e2e-test/finalizer_test.go
index f91ca52687..710b711676 100644
--- a/cmd/e2e-test/finalizer_test.go
+++ b/cmd/e2e-test/finalizer_test.go
@@ -58,9 +58,8 @@ func TestFinalizer(t *testing.T) {
 			t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
 		}
 
-		ingFinalizers := ing.GetFinalizers()
-		if len(ingFinalizers) != 1 || ingFinalizers[0] != common.FinalizerKey {
-			t.Fatalf("GetFinalizers() = %+v, want [%q]", ingFinalizers, common.FinalizerKey)
+		if err := e2e.CheckForAnyFinalizer(ing); err != nil {
+			t.Fatalf("CheckForAnyFinalizer(%s) = %v, want nil", ingKey, err)
 		}
 
 		// Perform whitebox testing.
@@ -109,9 +108,8 @@ func TestFinalizerIngressClassChange(t *testing.T) {
 			t.Fatalf("error waiting for Ingress %s to stabilize: %v", ingKey, err)
 		}
 
-		ingFinalizers := ing.GetFinalizers()
-		if len(ingFinalizers) != 1 || ingFinalizers[0] != common.FinalizerKey {
-			t.Fatalf("GetFinalizers() = %+v, want [%q]", ingFinalizers, common.FinalizerKey)
+		if err := e2e.CheckForAnyFinalizer(ing); err != nil {
+			t.Fatalf("CheckForAnyFinalizer(%s) = %v, want nil", ingKey, err)
 		}
 
 		// Perform whitebox testing.
@@ -184,13 +182,11 @@ func TestFinalizerIngressesWithSharedResources(t *testing.T) {
 			t.Fatalf("error waiting for Ingress %s to stabilize: %v", otherIngKey, err)
 		}
 
-		ingFinalizers := ing.GetFinalizers()
-		if len(ingFinalizers) != 1 || ingFinalizers[0] != common.FinalizerKey {
-			t.Fatalf("GetFinalizers() = %+v, want [%q]", ingFinalizers, common.FinalizerKey)
+		if err := e2e.CheckForAnyFinalizer(ing); err != nil {
+			t.Fatalf("CheckForAnyFinalizer(%s) = %v, want nil", ingKey, err)
 		}
-		otherIngFinalizers := otherIng.GetFinalizers()
-		if len(otherIngFinalizers) != 1 || otherIngFinalizers[0] != common.FinalizerKey {
-			t.Fatalf("GetFinalizers() = %+v, want [%q]", otherIngFinalizers, common.FinalizerKey)
+		if err := e2e.CheckForAnyFinalizer(otherIng); err != nil {
+			t.Fatalf("CheckForAnyFinalizer(%s) = %v, want nil", otherIngKey, err)
 		}
 
 		// Perform whitebox testing.
diff --git a/cmd/e2e-test/upgrade/finalizer.go b/cmd/e2e-test/upgrade/finalizer.go
index 549536b6ed..8904a98615 100644
--- a/cmd/e2e-test/upgrade/finalizer.go
+++ b/cmd/e2e-test/upgrade/finalizer.go
@@ -103,11 +103,15 @@ func (fr *Finalizer) PostUpgrade() error {
 	// Ingress status is updated to unstable, which would be set back to stable
 	// after WaitForFinalizer below finishes successfully.
 	fr.s.PutStatus(e2e.Unstable)
-	// Wait for finalizer to be added and verify that correct finalizer is added to the ingress after the upgrade.
+	// Wait for an ingress finalizer to be added on the ingress after the upgrade.
 	ing, err := e2e.WaitForFinalizer(fr.s, fr.ing)
 	if err != nil {
 		fr.t.Fatalf("e2e.WaitForFinalizer(_, %q) = _, %v, want nil", ingKey, err)
 	}
+	// Assert that v1 finalizer is added.
+	if err := e2e.CheckV1Finalizer(ing); err != nil {
+		fr.t.Fatalf("CheckV1Finalizer(%s) = %v, want nil", ingKey, err)
+	}
 	gclb, err := e2e.WhiteboxTest(ing, fr.s, fr.framework.Cloud, "")
 	if err != nil {
 		fr.t.Fatalf("e2e.WhiteboxTest(%s, ...) = %v, want nil", ingKey, err)
diff --git a/pkg/e2e/helpers.go b/pkg/e2e/helpers.go
index 2fd247c84b..33a7936285 100644
--- a/pkg/e2e/helpers.go
+++ b/pkg/e2e/helpers.go
@@ -143,7 +143,9 @@ func WaitForFinalizer(s *Sandbox, ing *v1beta1.Ingress) (*v1beta1.Ingress, error
 			return false, nil
 		}
 		ingFinalizers := ing.GetFinalizers()
-		if len(ingFinalizers) != 1 || ingFinalizers[0] != common.FinalizerKey {
+		if l := len(ingFinalizers); l > 1 {
+			return true, fmt.Errorf("unexpected number of finalizers for ingress %v, got %d", ing, l)
+		} else if l != 1 {
 			klog.Infof("WaitForFinalizer(%s) = %v, finalizer not added for Ingress %v", ingKey, ingFinalizers, ing)
 			return false, nil
 		}
@@ -473,3 +475,27 @@ func CheckNegStatus(svc *v1.Service, expectSvcPors []string) (annotations.NegSta
 	}
 	return negStatus, nil
 }
+
+// CheckForAnyFinalizer asserts that an ingress finalizer exists on Ingress.
+func CheckForAnyFinalizer(ing *v1beta1.Ingress) error {
+	ingFinalizers := ing.GetFinalizers()
+	if l := len(ingFinalizers); l != 1 {
+		return fmt.Errorf("expected 1 Finalizer but got %d", l)
+	}
+	if ingFinalizers[0] != common.FinalizerKey && ingFinalizers[0] != common.FinalizerKeyV2 {
+		return fmt.Errorf("unexpected finalizer %q found", ingFinalizers[0])
+	}
+	return nil
+}
+
+// CheckV1Finalizer asserts that v1 finalizer exists on Ingress.
+func CheckV1Finalizer(ing *v1beta1.Ingress) error {
+	ingFinalizers := ing.GetFinalizers()
+	if l := len(ingFinalizers); l != 1 {
+		return fmt.Errorf("expected 1 Finalizer but got %d", l)
+	}
+	if ingFinalizers[0] != common.FinalizerKey {
+		return fmt.Errorf("expected Finalizer %q but got %q", common.FinalizerKey, ingFinalizers[0])
+	}
+	return nil
+}