Return concrete types from mock dgxa100 server instead of interfaces

This way the callers can extend these types to futher override their functions if desired, while still being able to assign them to the appropriate nvml interfaces. Signed-off-by: Kevin Klues <[email protected]>
NVIDIA · Apr 23, 2024 · 5e1cdb1 · 5e1cdb1
1 parent 93fa13d
commit 5e1cdb1
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/pkg/nvml/mock/dgxa100/dgxa100.go b/pkg/nvml/mock/dgxa100/dgxa100.go
@@ -70,7 +70,7 @@ var _ nvml.Device = (*Device)(nil)
 var _ nvml.GpuInstance = (*GpuInstance)(nil)
 var _ nvml.ComputeInstance = (*ComputeInstance)(nil)
 
-func New() nvml.Interface {
+func New() *Server {
 	server := &Server{
 		Devices: [8]nvml.Device{
 			NewDevice(0),
@@ -90,7 +90,7 @@ func New() nvml.Interface {
 	return server
 }
 
-func NewDevice(index int) nvml.Device {
+func NewDevice(index int) *Device {
 	device := &Device{
 		UUID:         "GPU-" + uuid.New().String(),
 		Name:         "Mock NVIDIA A100-SXM4-40GB",
@@ -111,7 +111,7 @@ func NewDevice(index int) nvml.Device {
 	return device
 }
 
-func NewGpuInstance(info nvml.GpuInstanceInfo) nvml.GpuInstance {
+func NewGpuInstance(info nvml.GpuInstanceInfo) *GpuInstance {
 	gi := &GpuInstance{
 		Info:                   info,
 		ComputeInstances:       make(map[*ComputeInstance]struct{}),
@@ -121,7 +121,7 @@ func NewGpuInstance(info nvml.GpuInstanceInfo) nvml.GpuInstance {
 	return gi
 }
 
-func NewComputeInstance(info nvml.ComputeInstanceInfo) nvml.ComputeInstance {
+func NewComputeInstance(info nvml.ComputeInstanceInfo) *ComputeInstance {
 	ci := &ComputeInstance{
 		Info: info,
 	}
@@ -261,7 +261,7 @@ func (d *Device) setMockFuncs() {
 		}
 		d.GpuInstanceCounter++
 		gi := NewGpuInstance(giInfo)
-		d.GpuInstances[gi.(*GpuInstance)] = struct{}{}
+		d.GpuInstances[gi] = struct{}{}
 		return gi, nvml.SUCCESS
 	}
 
@@ -274,7 +274,7 @@ func (d *Device) setMockFuncs() {
 		}
 		d.GpuInstanceCounter++
 		gi := NewGpuInstance(giInfo)
-		d.GpuInstances[gi.(*GpuInstance)] = struct{}{}
+		d.GpuInstances[gi] = struct{}{}
 		return gi, nvml.SUCCESS
 	}
 
@@ -329,7 +329,7 @@ func (gi *GpuInstance) setMockFuncs() {
 		}
 		gi.ComputeInstanceCounter++
 		ci := NewComputeInstance(ciInfo)
-		gi.ComputeInstances[ci.(*ComputeInstance)] = struct{}{}
+		gi.ComputeInstances[ci] = struct{}{}
 		return ci, nvml.SUCCESS
 	}