cgroups: Unify TestAddCommand

Convert the v1 and v2 versions of `TestAddCommand` to a single shared
test.

3 files changed, 63 insertions, 116 deletions

diff --git a/internal/cgroups/handler_linux_test.go b/internal/cgroups/handler_linux_test.go
index 5e5a03952..9e8b7395c 100644
--- a/internal/cgroups/handler_linux_test.go
+++ b/internal/cgroups/handler_linux_test.go
@@ -265,6 +265,69 @@ func TestRepoCgroups(t *testing.T) {
 	}
 }
 
+func TestAddCommand(t *testing.T) {
+	for _, version := range []int{1, 2} {
+		t.Run("cgroups-v"+strconv.Itoa(version), func(t *testing.T) {
+			mock := newMock(t, version)
+
+			config := defaultCgroupsConfig()
+			config.Repositories.Count = 10
+			config.Repositories.MemoryBytes = 1024
+			config.Repositories.CPUShares = 16
+			config.HierarchyRoot = "gitaly"
+			config.Mountpoint = mock.rootPath()
+
+			pid := 1
+			manager1 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
+			mock.setupMockCgroupFiles(t, manager1, []uint{})
+			require.NoError(t, manager1.Setup())
+
+			ctx := testhelper.Context(t)
+
+			cmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
+			require.NoError(t, cmd2.Run())
+
+			groupID := calcGroupID(cmd2.Args, config.Repositories.Count)
+
+			manager2 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
+
+			t.Run("without overridden key", func(t *testing.T) {
+				_, err := manager2.AddCommand(cmd2)
+				require.NoError(t, err)
+				requireShards(t, version, mock, manager2, pid, groupID)
+
+				for _, path := range mock.repoPaths(pid, groupID) {
+					procsPath := filepath.Join(path, "cgroup.procs")
+					content := readCgroupFile(t, procsPath)
+
+					cmdPid, err := strconv.Atoi(string(content))
+					require.NoError(t, err)
+
+					require.Equal(t, cmd2.Process.Pid, cmdPid)
+				}
+			})
+
+			t.Run("with overridden key", func(t *testing.T) {
+				overriddenGroupID := calcGroupID([]string{"foobar"}, config.Repositories.Count)
+
+				_, err := manager2.AddCommand(cmd2, WithCgroupKey("foobar"))
+				require.NoError(t, err)
+				requireShards(t, version, mock, manager2, pid, groupID, overriddenGroupID)
+
+				for _, path := range mock.repoPaths(pid, overriddenGroupID) {
+					procsPath := filepath.Join(path, "cgroup.procs")
+					content := readCgroupFile(t, procsPath)
+
+					cmdPid, err := strconv.Atoi(string(content))
+					require.NoError(t, err)
+
+					require.Equal(t, cmd2.Process.Pid, cmdPid)
+				}
+			})
+		})
+	}
+}
+
 func requireCgroupComponents(t *testing.T, version int, root string, cgroupPath string, expected expectedCgroup) {
 	t.Helper()
 
diff --git a/internal/cgroups/v1_linux_test.go b/internal/cgroups/v1_linux_test.go
index e8840054f..531b42aca 100644
--- a/internal/cgroups/v1_linux_test.go
+++ b/internal/cgroups/v1_linux_test.go
@@ -33,64 +33,6 @@ func defaultCgroupsConfig() cgroups.Config {
 	}
 }
 
-func TestAddCommand(t *testing.T) {
-	mock := newMockV1(t)
-
-	config := defaultCgroupsConfig()
-	config.Repositories.Count = 10
-	config.Repositories.MemoryBytes = 1024
-	config.Repositories.CPUShares = 16
-	config.HierarchyRoot = "gitaly"
-	config.Mountpoint = mock.root
-
-	pid := 1
-	groupID := calcGroupID(cmdArgs, config.Repositories.Count)
-
-	v1Manager1 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
-	require.NoError(t, v1Manager1.Setup())
-	ctx := testhelper.Context(t)
-
-	cmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
-	require.NoError(t, cmd2.Run())
-
-	v1Manager2 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
-
-	t.Run("without overridden key", func(t *testing.T) {
-		_, err := v1Manager2.AddCommand(cmd2)
-		require.NoError(t, err)
-		requireShardsV1(t, mock, v1Manager2, pid, groupID)
-
-		for _, s := range mock.subsystems {
-			path := filepath.Join(mock.root, string(s.Name()), "gitaly",
-				fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
-			content := readCgroupFile(t, path)
-
-			cmdPid, err := strconv.Atoi(string(content))
-			require.NoError(t, err)
-
-			require.Equal(t, cmd2.Process.Pid, cmdPid)
-		}
-	})
-
-	t.Run("with overridden key", func(t *testing.T) {
-		overridenGroupID := calcGroupID([]string{"foobar"}, config.Repositories.Count)
-
-		_, err := v1Manager2.AddCommand(cmd2, WithCgroupKey("foobar"))
-		require.NoError(t, err)
-
-		for _, s := range mock.subsystems {
-			path := filepath.Join(mock.root, string(s.Name()), "gitaly",
-				fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", overridenGroupID), "cgroup.procs")
-			content := readCgroupFile(t, path)
-
-			cmdPid, err := strconv.Atoi(string(content))
-			require.NoError(t, err)
-
-			require.Equal(t, cmd2.Process.Pid, cmdPid)
-		}
-	})
-}
-
 func TestCleanup(t *testing.T) {
 	mock := newMockV1(t)
 
diff --git a/internal/cgroups/v2_linux_test.go b/internal/cgroups/v2_linux_test.go
index d3a459f85..f9c5f7bee 100644
--- a/internal/cgroups/v2_linux_test.go
+++ b/internal/cgroups/v2_linux_test.go
@@ -33,64 +33,6 @@ func defaultCgroupsV2Config() cgroups.Config {
 	}
 }
 
-func TestAddCommandV2(t *testing.T) {
-	mock := newMockV2(t)
-
-	config := defaultCgroupsV2Config()
-	config.Repositories.Count = 10
-	config.Repositories.MemoryBytes = 1024
-	config.Repositories.CPUShares = 16
-	config.Mountpoint = mock.root
-
-	pid := 1
-	groupID := calcGroupID(cmdArgs, config.Repositories.Count)
-
-	v2Manager1 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
-	mock.setupMockCgroupFiles(t, v2Manager1, []uint{})
-
-	require.NoError(t, v2Manager1.Setup())
-	ctx := testhelper.Context(t)
-
-	cmd2 := exec.CommandContext(ctx, cmdArgs[0], cmdArgs[1:]...)
-	require.NoError(t, cmd2.Run())
-
-	v2Manager2 := mock.newCgroupManager(config, testhelper.SharedLogger(t), pid)
-
-	t.Run("without overridden key", func(t *testing.T) {
-		groupID := calcGroupID(cmd2.Args, config.Repositories.Count)
-
-		_, err := v2Manager2.AddCommand(cmd2)
-		require.NoError(t, err)
-		requireShardsV2(t, mock, v2Manager2, pid, groupID)
-
-		path := filepath.Join(mock.root, "gitaly",
-			fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", groupID), "cgroup.procs")
-		content := readCgroupFile(t, path)
-
-		cmdPid, err := strconv.Atoi(string(content))
-		require.NoError(t, err)
-
-		require.Equal(t, cmd2.Process.Pid, cmdPid)
-	})
-
-	t.Run("with overridden key", func(t *testing.T) {
-		overriddenGroupID := calcGroupID([]string{"foobar"}, config.Repositories.Count)
-
-		_, err := v2Manager2.AddCommand(cmd2, WithCgroupKey("foobar"))
-		require.NoError(t, err)
-		requireShardsV2(t, mock, v2Manager2, pid, groupID, overriddenGroupID)
-
-		path := filepath.Join(mock.root, "gitaly",
-			fmt.Sprintf("gitaly-%d", pid), fmt.Sprintf("repos-%d", overriddenGroupID), "cgroup.procs")
-		content := readCgroupFile(t, path)
-
-		cmdPid, err := strconv.Atoi(string(content))
-		require.NoError(t, err)
-
-		require.Equal(t, cmd2.Process.Pid, cmdPid)
-	})
-}
-
 func TestCleanupV2(t *testing.T) {
 	mock := newMockV2(t)