@@ -20,7 +20,8 @@ import (
2020 "testing"
2121
2222 "github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
23- "github.com/NVIDIA/go-nvlib/pkg/nvml"
23+ "github.com/NVIDIA/go-nvml/pkg/nvml"
24+ "github.com/NVIDIA/go-nvml/pkg/nvml/mock"
2425 "github.com/stretchr/testify/require"
2526)
2627
@@ -32,7 +33,7 @@ func TestUsesNVGPUModule(t *testing.T) {
3233 }{
3334 {
3435 description : "init failure returns false" ,
35- nvmllib : & nvml. InterfaceMock {
36+ nvmllib : & mock. Interface {
3637 InitFunc : func () nvml.Return {
3738 return nvml .ERROR_LIBRARY_NOT_FOUND
3839 },
@@ -41,7 +42,7 @@ func TestUsesNVGPUModule(t *testing.T) {
4142 },
4243 {
4344 description : "no devices returns false" ,
44- nvmllib : & nvml. InterfaceMock {
45+ nvmllib : & mock. Interface {
4546 InitFunc : func () nvml.Return {
4647 return nvml .SUCCESS
4748 },
@@ -56,7 +57,7 @@ func TestUsesNVGPUModule(t *testing.T) {
5657 },
5758 {
5859 description : "DeviceGetCount error returns false" ,
59- nvmllib : & nvml. InterfaceMock {
60+ nvmllib : & mock. Interface {
6061 InitFunc : func () nvml.Return {
6162 return nvml .SUCCESS
6263 },
@@ -71,7 +72,7 @@ func TestUsesNVGPUModule(t *testing.T) {
7172 },
7273 {
7374 description : "Failure to get device name returns false" ,
74- nvmllib : & nvml. InterfaceMock {
75+ nvmllib : & mock. Interface {
7576 InitFunc : func () nvml.Return {
7677 return nvml .SUCCESS
7778 },
@@ -82,7 +83,7 @@ func TestUsesNVGPUModule(t *testing.T) {
8283 return 1 , nvml .SUCCESS
8384 },
8485 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
85- device := & nvml. DeviceMock {
86+ device := & mock. Device {
8687 GetNameFunc : func () (string , nvml.Return ) {
8788 return "" , nvml .ERROR_UNKNOWN
8889 },
@@ -94,7 +95,7 @@ func TestUsesNVGPUModule(t *testing.T) {
9495 },
9596 {
9697 description : "nested panic returns false" ,
97- nvmllib : & nvml. InterfaceMock {
98+ nvmllib : & mock. Interface {
9899 InitFunc : func () nvml.Return {
99100 return nvml .SUCCESS
100101 },
@@ -105,7 +106,7 @@ func TestUsesNVGPUModule(t *testing.T) {
105106 return 1 , nvml .SUCCESS
106107 },
107108 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
108- device := & nvml. DeviceMock {
109+ device := & mock. Device {
109110 GetNameFunc : func () (string , nvml.Return ) {
110111 panic ("deep panic" )
111112 },
@@ -117,7 +118,7 @@ func TestUsesNVGPUModule(t *testing.T) {
117118 },
118119 {
119120 description : "Single device name with no nvgpu" ,
120- nvmllib : & nvml. InterfaceMock {
121+ nvmllib : & mock. Interface {
121122 InitFunc : func () nvml.Return {
122123 return nvml .SUCCESS
123124 },
@@ -128,7 +129,7 @@ func TestUsesNVGPUModule(t *testing.T) {
128129 return 1 , nvml .SUCCESS
129130 },
130131 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
131- device := & nvml. DeviceMock {
132+ device := & mock. Device {
132133 GetNameFunc : func () (string , nvml.Return ) {
133134 return "NVIDIA A100-SXM4-40GB" , nvml .SUCCESS
134135 },
@@ -140,7 +141,7 @@ func TestUsesNVGPUModule(t *testing.T) {
140141 },
141142 {
142143 description : "Single device name with nvgpu" ,
143- nvmllib : & nvml. InterfaceMock {
144+ nvmllib : & mock. Interface {
144145 InitFunc : func () nvml.Return {
145146 return nvml .SUCCESS
146147 },
@@ -151,7 +152,7 @@ func TestUsesNVGPUModule(t *testing.T) {
151152 return 1 , nvml .SUCCESS
152153 },
153154 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
154- device := & nvml. DeviceMock {
155+ device := & mock. Device {
155156 GetNameFunc : func () (string , nvml.Return ) {
156157 return "Orin (nvgpu)" , nvml .SUCCESS
157158 },
@@ -163,7 +164,7 @@ func TestUsesNVGPUModule(t *testing.T) {
163164 },
164165 {
165166 description : "Multiple device names with no nvgpu" ,
166- nvmllib : & nvml. InterfaceMock {
167+ nvmllib : & mock. Interface {
167168 InitFunc : func () nvml.Return {
168169 return nvml .SUCCESS
169170 },
@@ -174,7 +175,7 @@ func TestUsesNVGPUModule(t *testing.T) {
174175 return 2 , nvml .SUCCESS
175176 },
176177 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
177- device := & nvml. DeviceMock {
178+ device := & mock. Device {
178179 GetNameFunc : func () (string , nvml.Return ) {
179180 return "NVIDIA A100-SXM4-40GB" , nvml .SUCCESS
180181 },
@@ -186,7 +187,7 @@ func TestUsesNVGPUModule(t *testing.T) {
186187 },
187188 {
188189 description : "Multiple device names with nvgpu" ,
189- nvmllib : & nvml. InterfaceMock {
190+ nvmllib : & mock. Interface {
190191 InitFunc : func () nvml.Return {
191192 return nvml .SUCCESS
192193 },
@@ -197,7 +198,7 @@ func TestUsesNVGPUModule(t *testing.T) {
197198 return 2 , nvml .SUCCESS
198199 },
199200 DeviceGetHandleByIndexFunc : func (index int ) (nvml.Device , nvml.Return ) {
200- device := & nvml. DeviceMock {
201+ device := & mock. Device {
201202 GetNameFunc : func () (string , nvml.Return ) {
202203 return "Orin (nvgpu)" , nvml .SUCCESS
203204 },
@@ -209,7 +210,7 @@ func TestUsesNVGPUModule(t *testing.T) {
209210 },
210211 {
211212 description : "Mixed device names" ,
212- nvmllib : & nvml. InterfaceMock {
213+ nvmllib : & mock. Interface {
213214 InitFunc : func () nvml.Return {
214215 return nvml .SUCCESS
215216 },
@@ -226,7 +227,7 @@ func TestUsesNVGPUModule(t *testing.T) {
226227 } else {
227228 deviceName = "Orin (nvgpu)"
228229 }
229- device := & nvml. DeviceMock {
230+ device := & mock. Device {
230231 GetNameFunc : func () (string , nvml.Return ) {
231232 return deviceName , nvml .SUCCESS
232233 },
0 commit comments