From e5cbc51d29db751f074653fbbf5eeac2238834db Mon Sep 17 00:00:00 2001
From: Stanislav Laznicka <slaznick@redhat.com>
Date: Wed, 10 May 2023 15:25:50 +0200
Subject: [PATCH] e2e framework: allow setting all PSa labels at once

---
 test/e2e/framework/framework.go | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/test/e2e/framework/framework.go b/test/e2e/framework/framework.go
index ecd54705939..e29e0edf58e 100644
--- a/test/e2e/framework/framework.go
+++ b/test/e2e/framework/framework.go
@@ -89,6 +89,12 @@ var (
 
 // Framework supports common operations used by e2e tests; it will keep a client & a namespace for you.
 // Eventual goal is to merge this with integration test framework.
+//
+// You can configure the pod security level for your test by setting the `NamespacePodSecurityLevel`
+// which will set all three of pod security admission enforce, warn and audit labels on the namespace.
+// The default pod security profile is "restricted".
+// Each of the labels can be overridden by using more specific NamespacePodSecurity* attributes of this
+// struct.
 type Framework struct {
 	BaseName string
 
@@ -111,6 +117,9 @@ type Framework struct {
 	namespacesToDelete               []*v1.Namespace // Some tests have more than one.
 	NamespaceDeletionTimeout         time.Duration
 	NamespacePodSecurityEnforceLevel admissionapi.Level // The pod security enforcement level for namespaces to be applied.
+	NamespacePodSecurityWarnLevel    admissionapi.Level // The pod security warn (client logging) level for namespaces to be applied.
+	NamespacePodSecurityAuditLevel   admissionapi.Level // The pod security audit (server logging) level for namespaces to be applied.
+	NamespacePodSecurityLevel        admissionapi.Level // The pod security level to be used for all of enforcement, warn and audit. Can be rewritten by more specific configuration attributes.
 
 	// Flaky operation failures in an e2e test can be captured through this.
 	flakeReport *FlakeReport
@@ -448,11 +457,9 @@ func (f *Framework) CreateNamespace(ctx context.Context, baseName string, labels
 		labels = labelsCopy
 	}
 
-	enforceLevel := admissionapi.LevelRestricted
-	if f.NamespacePodSecurityEnforceLevel != "" {
-		enforceLevel = f.NamespacePodSecurityEnforceLevel
-	}
-	labels[admissionapi.EnforceLevelLabel] = string(enforceLevel)
+	labels[admissionapi.EnforceLevelLabel] = firstNonEmptyPSaLevelOrRestricted(f.NamespacePodSecurityEnforceLevel, f.NamespacePodSecurityLevel)
+	labels[admissionapi.WarnLevelLabel] = firstNonEmptyPSaLevelOrRestricted(f.NamespacePodSecurityWarnLevel, f.NamespacePodSecurityLevel)
+	labels[admissionapi.AuditLevelLabel] = firstNonEmptyPSaLevelOrRestricted(f.NamespacePodSecurityAuditLevel, f.NamespacePodSecurityLevel)
 
 	ns, err := createTestingNS(ctx, baseName, f.ClientSet, labels)
 	// check ns instead of err to see if it's nil as we may
@@ -481,6 +488,15 @@ func (f *Framework) CreateNamespace(ctx context.Context, baseName string, labels
 	return ns, err
 }
 
+func firstNonEmptyPSaLevelOrRestricted(levelConfig ...admissionapi.Level) string {
+	for _, l := range levelConfig {
+		if len(l) > 0 {
+			return string(l)
+		}
+	}
+	return string(admissionapi.LevelRestricted)
+}
+
 // createSecretFromDockerConfig creates a secret using the private image registry credentials.
 // The credentials are provided by --e2e-docker-config-file flag.
 func (f *Framework) createSecretFromDockerConfig(ctx context.Context, namespace string) (*v1.Secret, error) {