From 09fdae408fb1fefdae2a8eb38ca8b8f9c35483ed Mon Sep 17 00:00:00 2001
From: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
Date: Fri, 14 Feb 2025 10:47:03 +0900
Subject: [PATCH] pkg/kubelet/userns: add idsPerPod tests

Signed-off-by: Akihiro Suda <akihiro.suda.cz@hco.ntt.co.jp>
---
 pkg/kubelet/userns/userns_manager_test.go | 114 ++++++++++++++--------
 1 file changed, 72 insertions(+), 42 deletions(-)

diff --git a/pkg/kubelet/userns/userns_manager_test.go b/pkg/kubelet/userns/userns_manager_test.go
index 42899c763d5..cb66fb854e5 100644
--- a/pkg/kubelet/userns/userns_manager_test.go
+++ b/pkg/kubelet/userns/userns_manager_test.go
@@ -102,51 +102,81 @@ func (m *testUserNsPodsManager) GetUserNamespacesIDsPerPod() uint32 {
 func TestUserNsManagerAllocate(t *testing.T) {
 	featuregatetesting.SetFeatureGateDuringTest(t, utilfeature.DefaultFeatureGate, pkgfeatures.UserNamespacesSupport, true)
 
-	testUserNsPodsManager := &testUserNsPodsManager{}
-	m, err := MakeUserNsManager(testUserNsPodsManager)
-	require.NoError(t, err)
+	customUserNsLength := uint32(1048576)
 
-	allocated, length, err := m.allocateOne("one")
-	assert.NoError(t, err)
-	assert.Equal(t, testUserNsLength, length, "m.isSet(%d).length=%v", allocated, length)
-	assert.True(t, m.isSet(allocated), "m.isSet(%d)", allocated)
-
-	allocated2, length2, err := m.allocateOne("two")
-	assert.NoError(t, err)
-	assert.NotEqual(t, allocated, allocated2, "allocated != allocated2")
-	assert.Equal(t, length, length2, "length == length2")
-
-	// verify that re-adding the same pod with the same settings won't fail
-	err = m.record("two", allocated2, length2)
-	assert.NoError(t, err)
-	// but it fails if anyting is different
-	err = m.record("two", allocated2+1, length2)
-	assert.Error(t, err)
-
-	m.Release("one")
-	m.Release("two")
-	assert.False(t, m.isSet(allocated), "m.isSet(%d)", allocated)
-	assert.False(t, m.isSet(allocated2), "m.nsSet(%d)", allocated2)
-
-	var allocs []uint32
-	for i := 0; i < 1000; i++ {
-		allocated, length, err = m.allocateOne(types.UID(fmt.Sprintf("%d", i)))
-		assert.Equal(t, testUserNsLength, length, "length is not the expected. iter: %v", i)
-		assert.NoError(t, err)
-		assert.GreaterOrEqual(t, allocated, uint32(minimumMappingUID))
-		// The last ID of the userns range (allocated+userNsLength) should be within bounds.
-		assert.LessOrEqual(t, allocated, uint32(minimumMappingUID+mappingLen-testUserNsLength))
-		allocs = append(allocs, allocated)
+	cases := []struct {
+		name           string
+		userNsLength   uint32
+		mappingFirstID uint32
+		mappingLen     uint32
+	}{
+		{
+			name:           "default",
+			userNsLength:   testUserNsLength,
+			mappingFirstID: minimumMappingUID,
+			mappingLen:     mappingLen,
+		},
+		{
+			name:           "custom",
+			userNsLength:   customUserNsLength,
+			mappingFirstID: customUserNsLength,
+			mappingLen:     customUserNsLength * 2000,
+		},
 	}
-	for i, v := range allocs {
-		assert.True(t, m.isSet(v), "m.isSet(%d) should be true", v)
-		m.Release(types.UID(fmt.Sprintf("%d", i)))
-		assert.False(t, m.isSet(v), "m.isSet(%d) should be false", v)
 
-		err = m.record(types.UID(fmt.Sprintf("%d", i)), v, testUserNsLength)
-		assert.NoError(t, err)
-		m.Release(types.UID(fmt.Sprintf("%d", i)))
-		assert.False(t, m.isSet(v), "m.isSet(%d) should be false", v)
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			testUserNsPodsManager := &testUserNsPodsManager{
+				userNsLength:   tc.userNsLength,
+				mappingFirstID: tc.mappingFirstID,
+				mappingLen:     tc.mappingLen,
+			}
+			m, err := MakeUserNsManager(testUserNsPodsManager)
+			require.NoError(t, err)
+
+			allocated, length, err := m.allocateOne("one")
+			require.NoError(t, err)
+			assert.Equal(t, tc.userNsLength, length, "m.isSet(%d).length=%v", allocated, length)
+			assert.True(t, m.isSet(allocated), "m.isSet(%d)", allocated)
+
+			allocated2, length2, err := m.allocateOne("two")
+			require.NoError(t, err)
+			assert.NotEqual(t, allocated, allocated2, "allocated != allocated2")
+			assert.Equal(t, length, length2, "length == length2")
+
+			// verify that re-adding the same pod with the same settings won't fail
+			err = m.record("two", allocated2, length2)
+			require.NoError(t, err)
+			// but it fails if anyting is different
+			err = m.record("two", allocated2+1, length2)
+			require.Error(t, err)
+
+			m.Release("one")
+			m.Release("two")
+			assert.False(t, m.isSet(allocated), "m.isSet(%d)", allocated)
+			assert.False(t, m.isSet(allocated2), "m.nsSet(%d)", allocated2)
+
+			var allocs []uint32
+			for i := 0; i < 1000; i++ {
+				allocated, length, err = m.allocateOne(types.UID(fmt.Sprintf("%d", i)))
+				assert.Equal(t, tc.userNsLength, length, "length is not the expected. iter: %v", i)
+				require.NoError(t, err)
+				assert.GreaterOrEqual(t, allocated, tc.mappingFirstID)
+				// The last ID of the userns range (allocated+userNsLength) should be within bounds.
+				assert.LessOrEqual(t, allocated, tc.mappingFirstID+tc.mappingLen-tc.userNsLength)
+				allocs = append(allocs, allocated)
+			}
+			for i, v := range allocs {
+				assert.True(t, m.isSet(v), "m.isSet(%d) should be true", v)
+				m.Release(types.UID(fmt.Sprintf("%d", i)))
+				assert.False(t, m.isSet(v), "m.isSet(%d) should be false", v)
+
+				err = m.record(types.UID(fmt.Sprintf("%d", i)), v, tc.userNsLength)
+				require.NoError(t, err)
+				m.Release(types.UID(fmt.Sprintf("%d", i)))
+				assert.False(t, m.isSet(v), "m.isSet(%d) should be false", v)
+			}
+		})
 	}
 }