Add ut coverage for capabilities.Setup (#125395)

* Add ut coverage for capabilities.Setup

* Update pkg/capabilities/capabilities_test.go

Co-authored-by: Ed Bartosh <eduard.bartosh@intel.com>

* Add ut coverage for capabilities.Setup

Signed-off-by: robert-cronin <robert.owen.cronin@gmail.com>

---------

Signed-off-by: robert-cronin <robert.owen.cronin@gmail.com>
Co-authored-by: Ed Bartosh <eduard.bartosh@intel.com>

pkg/apis/core/validation/validation_test.go

@@ -8323,7 +8323,8 @@ func TestValidateLinuxPodSecurityContext(t *testing.T) {
 
 func TestValidateContainers(t *testing.T) {
 	volumeDevices := make(map[string]core.VolumeSource)
-	capabilities.SetForTests(capabilities.Capabilities{
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{
 		AllowPrivileged: true,
 	})
 
@@ -8526,7 +8527,8 @@ func TestValidateContainers(t *testing.T) {
 		t.Errorf("expected success: %v", errs)
 	}
 
-	capabilities.SetForTests(capabilities.Capabilities{
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{
 		AllowPrivileged: false,
 	})
 	errorCases := []struct {
@@ -9151,7 +9153,8 @@ func TestValidateContainers(t *testing.T) {
 
 func TestValidateInitContainers(t *testing.T) {
 	volumeDevices := make(map[string]core.VolumeSource)
-	capabilities.SetForTests(capabilities.Capabilities{
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{
 		AllowPrivileged: true,
 	})
 
@@ -9229,7 +9232,8 @@ func TestValidateInitContainers(t *testing.T) {
 		t.Errorf("expected success: %v", errs)
 	}
 
-	capabilities.SetForTests(capabilities.Capabilities{
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{
 		AllowPrivileged: false,
 	})
 	errorCases := []struct {
@@ -14508,7 +14512,8 @@ func TestValidatePodEphemeralContainersUpdate(t *testing.T) {
 
 	// Some tests use Windows host pods as an example of fields that might
 	// conflict between an ephemeral container and the rest of the pod.
-	capabilities.SetForTests(capabilities.Capabilities{
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{
 		AllowPrivileged: true,
 	})
 	makeWindowsHostPod := func(ephemeralContainers []core.EphemeralContainer) *core.Pod {
@@ -20996,7 +21001,8 @@ func TestValidateSecurityContext(t *testing.T) {
 		},
 	}
 	for k, v := range errorCases {
-		capabilities.SetForTests(capabilities.Capabilities{
+		capabilities.ResetForTest()
+		capabilities.Initialize(capabilities.Capabilities{
 			AllowPrivileged: v.capAllowPriv,
 		})
 		// note the unconditional `true` here for hostUsers. The failure case to test for ProcMount only includes it being true,
@@ -23649,8 +23655,8 @@ func TestValidateWindowsHostProcessPod(t *testing.T) {
 
 	for _, testCase := range testCases {
 		t.Run(testCase.name, func(t *testing.T) {
-
+			capabilities.ResetForTest()
-			capabilities.SetForTests(capabilities.Capabilities{
+			capabilities.Initialize(capabilities.Capabilities{
 				AllowPrivileged: testCase.allowPrivileged,
 			})
 

pkg/capabilities/capabilities.go

@@ -68,11 +68,13 @@ func Setup(allowPrivileged bool, perConnectionBytesPerSec int64) {
 	})
 }
 
-// SetForTests sets capabilities for tests.  Convenience method for testing.  This should only be called from tests.
+// ResetForTest resets the capabilities to a given state for testing purposes.
-func SetForTests(c Capabilities) {
+// This function should only be called from tests.
+func ResetForTest() {
 	capInstance.lock.Lock()
 	defer capInstance.lock.Unlock()
-	capInstance.capabilities = &c
+	capInstance.capabilities = nil
+	capInstance.once = sync.Once{}
 }
 
 // Get returns a read-only copy of the system capabilities.

pkg/capabilities/capabilities_test.go

@@ -18,17 +18,11 @@ package capabilities
 
 import (
 	"reflect"
-	"sync"
 	"testing"
 )
 
 func TestGet(t *testing.T) {
-	defer func() {
+	defer ResetForTest()
-		capInstance.lock.Lock()
-		defer capInstance.lock.Unlock()
-		capInstance.capabilities = nil
-		capInstance.once = sync.Once{}
-	}()
 	defaultCap := Capabilities{
 		AllowPrivileged: false,
 		PrivilegedSources: PrivilegedSources{
@@ -48,10 +42,51 @@ func TestGet(t *testing.T) {
 			HostNetworkSources: []string{"A", "B"},
 		},
 	}
-	SetForTests(cap)
+	ResetForTest()
+	Initialize(cap)
 
 	res = Get()
 	if !reflect.DeepEqual(cap, res) {
 		t.Fatalf("expected Capabilities: %#v , got a different: %#v", cap, res)
 	}
 }
+func TestSetup(t *testing.T) {
+	defer ResetForTest()
+	testCases := []struct {
+		name                     string
+		allowPrivileged          bool
+		perConnectionBytesPerSec int64
+		expectedCapabilities     Capabilities
+	}{
+		{
+			name:                     "AllowPrivileged true with bandwidth limit",
+			allowPrivileged:          true,
+			perConnectionBytesPerSec: 1024,
+			expectedCapabilities: Capabilities{
+				AllowPrivileged:                        true,
+				PerConnectionBandwidthLimitBytesPerSec: 1024,
+			},
+		},
+		{
+			name:                     "AllowPrivileged false with higher bandwidth limit",
+			allowPrivileged:          false,
+			perConnectionBytesPerSec: 2048,
+			expectedCapabilities: Capabilities{
+				AllowPrivileged:                        false,
+				PerConnectionBandwidthLimitBytesPerSec: 2048,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ResetForTest()
+
+			Setup(tc.allowPrivileged, tc.perConnectionBytesPerSec)
+			res := Get()
+			if !reflect.DeepEqual(tc.expectedCapabilities, res) {
+				t.Fatalf("expected Capabilities: %#v, got: %#v", tc.expectedCapabilities, res)
+			}
+		})
+	}
+}

test/integration/auth/podsecurity_test.go

@@ -102,7 +102,8 @@ func TestPodSecurityWebhook(t *testing.T) {
 	featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, features.UserNamespacesSupport, true)
 
 	// Start test API server.
-	capabilities.SetForTests(capabilities.Capabilities{AllowPrivileged: true})
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{AllowPrivileged: true})
 	testServer := kubeapiservertesting.StartTestServerOrDie(t, kubeapiservertesting.NewDefaultTestServerOptions(), []string{
 		"--anonymous-auth=false",
 		"--allow-privileged=true",
@@ -136,7 +137,8 @@ func TestPodSecurityWebhook(t *testing.T) {
 
 func startPodSecurityServer(t *testing.T) *kubeapiservertesting.TestServer {
 	// ensure the global is set to allow privileged containers
-	capabilities.SetForTests(capabilities.Capabilities{AllowPrivileged: true})
+	capabilities.ResetForTest()
+	capabilities.Initialize(capabilities.Capabilities{AllowPrivileged: true})
 
 	server := kubeapiservertesting.StartTestServerOrDie(t, kubeapiservertesting.NewDefaultTestServerOptions(), []string{
 		"--anonymous-auth=false",