Add cleanup to reset security group on aigrapped registry mirror dock…

…er tests (#7245)
aws · Jan 2, 2024 · d349dc5 · d349dc5
1 parent 48435f1
commit d349dc5
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/test/e2e/airgap.go b/test/e2e/airgap.go
@@ -189,13 +189,14 @@ func runDockerAirgapConfigFlow(test *framework.ClusterE2ETest) {
 	test.ExtractDownloadedArtifacts()
 	test.DownloadImages()
 	test.ChangeInstanceSecurityGroup(os.Getenv(framework.RegistryMirrorAirgappedSecurityGroup))
+	test.SetRegistryMirrorDefaultInstanceSecurityGroupOnCleanup()
 	test.ImportImages()
 	test.CreateCluster(framework.WithBundlesOverride(bundleReleasePathFromArtifacts))
 	test.DeleteCluster(framework.WithBundlesOverride(bundleReleasePathFromArtifacts))
-	test.ChangeInstanceSecurityGroup(os.Getenv(framework.RegistryMirrorDefaultSecurityGroup))
 }
 
 func runDockerAirgapUpgradeFromReleaseFlow(test *framework.ClusterE2ETest, latestRelease *releasev1.EksARelease, wantVersion anywherev1.KubernetesVersion) {
+	test.SetRegistryMirrorDefaultInstanceSecurityGroupOnCleanup()
 	test.GenerateClusterConfigForVersion(latestRelease.Version, framework.ExecuteWithEksaRelease(latestRelease))
 
 	// Downloading and importing the artifacts from the previous version
@@ -214,11 +215,11 @@ func runDockerAirgapUpgradeFromReleaseFlow(test *framework.ClusterE2ETest, lates
 	test.ExtractDownloadedArtifacts()
 	test.DownloadImages()
 	test.ChangeInstanceSecurityGroup(os.Getenv(framework.RegistryMirrorAirgappedSecurityGroup))
+	test.SetRegistryMirrorDefaultInstanceSecurityGroupOnCleanup()
 	test.ImportImages()
 
 	test.UpgradeClusterWithNewConfig([]framework.ClusterE2ETestOpt{}, framework.WithBundlesOverride(bundleReleasePathFromArtifacts))
 	test.ValidateCluster(wantVersion)
 	test.StopIfFailed()
 	test.DeleteCluster(framework.WithBundlesOverride(bundleReleasePathFromArtifacts))
-	test.ChangeInstanceSecurityGroup(os.Getenv(framework.RegistryMirrorDefaultSecurityGroup))
 }
diff --git a/test/framework/registry_mirror.go b/test/framework/registry_mirror.go
@@ -164,3 +164,10 @@ func setupRegistryMirrorEndpointAndCert(e *ClusterE2ETest, providerName string,
 		e.T.Fatalf("unable to set REGISTRY_PASSWORD: %v", err)
 	}
 }
+
+// SetRegistryMirrorDefaultInstanceSecurityGroupOnCleanup sets the instance security group to the registry mirror default security group on cleanup.
+func (e *ClusterE2ETest) SetRegistryMirrorDefaultInstanceSecurityGroupOnCleanup(opts ...CommandOpt) {
+	e.T.Cleanup(func() {
+		e.ChangeInstanceSecurityGroup(os.Getenv(RegistryMirrorDefaultSecurityGroup))
+	})
+}