ClearML/go-nvlib
1
0
Fork 0
mirror of https://github.com/clearml/go-nvlib synced 2025-03-04 02:57:44 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Extract MockNVDeviceLib into helper function in nvdev tests

Browse Source 
Signed-off-by: Kevin Klues <kklues@nvidia.com>
This commit is contained in:
Kevin Klues 2023-03-26 13:45:05 +00:00
parent bdc2553e78
commit 2b4f40a90b
1 changed files with 53 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

102
pkg/nvlib/device/mig_profile_test.go
View File

@ -24,6 +24,58 @@ import (
"gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml" "gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml"
) )
func newMockDeviceLib() Interface {
mockDevice := &nvml.DeviceMock{
GetNameFunc: func() (string, nvml.Return) {
return "MockDevice", nvml.SUCCESS
},
GetMigModeFunc: func() (int, int, nvml.Return) {
return nvml.DEVICE_MIG_ENABLE, nvml.DEVICE_MIG_ENABLE, nvml.SUCCESS
},
GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) {
memory := nvml.Memory{
Total: 40 * 1024 * 1024 * 1024,
}
return memory, nvml.SUCCESS
},
GetGpuInstanceProfileInfoFunc: func(Profile int) (nvml.GpuInstanceProfileInfo, nvml.Return) {
info := nvml.GpuInstanceProfileInfo{}
switch Profile {
case nvml.GPU_INSTANCE_PROFILE_1_SLICE,
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1:
info.MemorySizeMB = 5 * 1024
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV2:
info.MemorySizeMB = 10 * 1024
case nvml.GPU_INSTANCE_PROFILE_2_SLICE,
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1:
info.MemorySizeMB = 10 * 1024
case nvml.GPU_INSTANCE_PROFILE_3_SLICE:
info.MemorySizeMB = 20 * 1024
case nvml.GPU_INSTANCE_PROFILE_4_SLICE:
info.MemorySizeMB = 20 * 1024
case nvml.GPU_INSTANCE_PROFILE_7_SLICE:
info.MemorySizeMB = 40 * 1024
case nvml.GPU_INSTANCE_PROFILE_6_SLICE,
nvml.GPU_INSTANCE_PROFILE_8_SLICE:
fallthrough
default:
return info, nvml.ERROR_NOT_SUPPORTED
}
return info, nvml.SUCCESS
},
}
mockNvml := &nvml.InterfaceMock{
DeviceGetCountFunc: func() (int, nvml.Return) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(Index int) (nvml.Device, nvml.Return) {
return mockDevice, nvml.SUCCESS
},
}
return New(WithNvml(mockNvml), WithVerifySymbols(false))
}
func TestParseMigProfile(t *testing.T) { func TestParseMigProfile(t *testing.T) {
testCases := []struct { testCases := []struct {
description string description string
@ -291,55 +343,7 @@ func TestParseMigProfile(t *testing.T) {
}, },
} }
mockDevice := &nvml.DeviceMock{ d := newMockDeviceLib()
GetNameFunc: func() (string, nvml.Return) {
return "MockDevice", nvml.SUCCESS
},
GetMigModeFunc: func() (int, int, nvml.Return) {
return nvml.DEVICE_MIG_ENABLE, nvml.DEVICE_MIG_ENABLE, nvml.SUCCESS
},
GetMemoryInfoFunc: func() (nvml.Memory, nvml.Return) {
memory := nvml.Memory{
Total: 40 * 1024 * 1024 * 1024,
}
return memory, nvml.SUCCESS
},
GetGpuInstanceProfileInfoFunc: func(Profile int) (nvml.GpuInstanceProfileInfo, nvml.Return) {
info := nvml.GpuInstanceProfileInfo{}
switch Profile {
case nvml.GPU_INSTANCE_PROFILE_1_SLICE,
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1:
info.MemorySizeMB = 5 * 1024
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV2:
info.MemorySizeMB = 10 * 1024
case nvml.GPU_INSTANCE_PROFILE_2_SLICE,
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1:
info.MemorySizeMB = 10 * 1024
case nvml.GPU_INSTANCE_PROFILE_3_SLICE:
info.MemorySizeMB = 20 * 1024
case nvml.GPU_INSTANCE_PROFILE_4_SLICE:
info.MemorySizeMB = 20 * 1024
case nvml.GPU_INSTANCE_PROFILE_7_SLICE:
info.MemorySizeMB = 40 * 1024
case nvml.GPU_INSTANCE_PROFILE_6_SLICE,
nvml.GPU_INSTANCE_PROFILE_8_SLICE:
fallthrough
default:
return info, nvml.ERROR_NOT_SUPPORTED
}
return info, nvml.SUCCESS
},
}
mockNvml := &nvml.InterfaceMock{
DeviceGetCountFunc: func() (int, nvml.Return) {
return 1, nvml.SUCCESS
},
DeviceGetHandleByIndexFunc: func(Index int) (nvml.Device, nvml.Return) {
return mockDevice, nvml.SUCCESS
},
}
d := New(WithNvml(mockNvml), WithVerifySymbols(false))
for _, tc := range testCases { for _, tc := range testCases {
t.Run(tc.description, func(t *testing.T) { t.Run(tc.description, func(t *testing.T) {
err := d.AssertValidMigProfileFormat(tc.device) err := d.AssertValidMigProfileFormat(tc.device)