From ca399c46bb03a14f550e8d837e5dd4d5018fc9b6 Mon Sep 17 00:00:00 2001
From: Kevin Klues <kklues@nvidia.com>
Date: Fri, 12 Apr 2024 21:26:47 +0000
Subject: [PATCH] Parallelize tests to allow us to reenable testing in pkg/...

Signed-off-by: Kevin Klues <kklues@nvidia.com>
---
 Makefile                             |  2 +-
 pkg/mig/config/config_test.go        | 31 +++++++++++++++++++---------
 pkg/mig/config/known_configs_test.go |  4 +++-
 pkg/mig/state/state_test.go          | 18 +++++++++-------
 4 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/Makefile b/Makefile
index dd8ef674..af5b121e 100644
--- a/Makefile
+++ b/Makefile
@@ -83,7 +83,7 @@ check-vendor: vendor
 
 COVERAGE_FILE := coverage.out
 test: build cmds
-	go test -coverprofile=$(COVERAGE_FILE) $(MODULE)/cmd/... $(MODULE)/internal/... $(MODULE)/api/...
+	go test -v -coverprofile=$(COVERAGE_FILE) $(MODULE)/cmd/... $(MODULE)/internal/... $(MODULE)/api/... $(MODULE)/pkg/...
 
 coverage: test
 	cat $(COVERAGE_FILE) | grep -v "_mock.go" > $(COVERAGE_FILE).no-mocks
diff --git a/pkg/mig/config/config_test.go b/pkg/mig/config/config_test.go
index ea326378..5cf9cc5b 100644
--- a/pkg/mig/config/config_test.go
+++ b/pkg/mig/config/config_test.go
@@ -44,13 +44,7 @@ func EnableMigMode(manager Manager, gpu int) (nvml.Return, nvml.Return) {
 }
 
 func TestGetSetMigConfig(t *testing.T) {
-	nvmlLib := dgxa100.New()
-	manager := NewMockLunaServerMigConfigManager()
-
-	numGPUs, ret := nvmlLib.DeviceGetCount()
-	require.NotNil(t, ret, "Unexpected nil return from DeviceGetCount")
-	require.Equal(t, ret, nvml.SUCCESS, "Unexpected return value from DeviceGetCount")
-
+	types.SetMockNVdevlib()
 	mcg := NewA100_SXM4_40GB_MigConfigGroup()
 
 	type testCase struct {
@@ -69,8 +63,18 @@ func TestGetSetMigConfig(t *testing.T) {
 		return testCases
 	}()
 
-	for _, tc := range testCases {
+	for i := range testCases {
+		tc := testCases[i] // to allow us to run parallelly
 		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
+			nvmlLib := dgxa100.New()
+			manager := NewMockLunaServerMigConfigManager()
+
+			numGPUs, ret := nvmlLib.DeviceGetCount()
+			require.NotNil(t, ret, "Unexpected nil return from DeviceGetCount")
+			require.Equal(t, ret, nvml.SUCCESS, "Unexpected return value from DeviceGetCount")
+
 			for i := 0; i < numGPUs; i++ {
 				r1, r2 := EnableMigMode(manager, i)
 				require.Equal(t, nvml.SUCCESS, r1)
@@ -88,6 +92,7 @@ func TestGetSetMigConfig(t *testing.T) {
 }
 
 func TestClearMigConfig(t *testing.T) {
+	types.SetMockNVdevlib()
 	mcg := NewA100_SXM4_40GB_MigConfigGroup()
 
 	type testCase struct {
@@ -106,8 +111,11 @@ func TestClearMigConfig(t *testing.T) {
 		return testCases
 	}()
 
-	for _, tc := range testCases {
+	for i := range testCases {
+		tc := testCases[i] // to allow us to run parallelly
 		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
 			manager := NewMockLunaServerMigConfigManager()
 
 			r1, r2 := EnableMigMode(manager, 0)
@@ -173,8 +181,11 @@ func TestIteratePermutationsUntilSuccess(t *testing.T) {
 		return testCases
 	}()
 
-	for _, tc := range testCases {
+	for i := range testCases {
+		tc := testCases[i] // to allow us to run parallelly
 		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
 			iteration := 0
 			err := iteratePermutationsUntilSuccess(tc.config, func(perm []*types.MigProfile) error {
 				iteration++
diff --git a/pkg/mig/config/known_configs_test.go b/pkg/mig/config/known_configs_test.go
index 33695ff5..bb337017 100644
--- a/pkg/mig/config/known_configs_test.go
+++ b/pkg/mig/config/known_configs_test.go
@@ -132,8 +132,10 @@ func TestValidConfiguration(t *testing.T) {
 	types.SetMockNVdevlib()
 	configs := GetKnownMigConfigGroups()
 
-	for _, tc := range testCases {
+	for i := range testCases {
+		tc := testCases[i] // to allow us to run parallelly
 		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
 			err := configs[tc.gpu].AssertValidConfiguration(tc.config)
 			if tc.valid {
 				require.Nil(t, err)
diff --git a/pkg/mig/state/state_test.go b/pkg/mig/state/state_test.go
index e9c161b3..f0709b67 100644
--- a/pkg/mig/state/state_test.go
+++ b/pkg/mig/state/state_test.go
@@ -36,12 +36,7 @@ func newMockMigStateManagerOnLunaServer() *migStateManager {
 }
 
 func TestFetchRestore(t *testing.T) {
-	manager := newMockMigStateManagerOnLunaServer()
-
-	numGPUs, ret := manager.nvml.DeviceGetCount()
-	require.NotNil(t, ret, "Unexpected nil return from DeviceGetCount")
-	require.Equal(t, ret, nvml.SUCCESS, "Unexpected return value from DeviceGetCount")
-
+	types.SetMockNVdevlib()
 	mcg := config.NewA100_SXM4_40GB_MigConfigGroup()
 
 	type testCase struct {
@@ -73,8 +68,17 @@ func TestFetchRestore(t *testing.T) {
 		return testCases
 	}()
 
-	for _, tc := range testCases {
+	for i := range testCases {
+		tc := testCases[i] // to allow us to run parallelly
 		t.Run(tc.description, func(t *testing.T) {
+			t.Parallel()
+
+			manager := newMockMigStateManagerOnLunaServer()
+
+			numGPUs, ret := manager.nvml.DeviceGetCount()
+			require.NotNil(t, ret, "Unexpected nil return from DeviceGetCount")
+			require.Equal(t, ret, nvml.SUCCESS, "Unexpected return value from DeviceGetCount")
+
 			for i := 0; i < numGPUs; i++ {
 				err := manager.mode.SetMigMode(i, tc.mode)
 				require.Nil(t, err)