Merge pull request #23 from elezar/fix-devicelist-filter

Fix devicelist filter

Author: elezar

go.mod

@@ -2,11 +2,17 @@ module github.com/NVIDIA/go-gpuallocator
 
 go 1.20
 
-require github.com/NVIDIA/go-nvlib v0.0.0-20240109130712-11603560817a
+require (
+	github.com/NVIDIA/go-nvlib v0.0.0-20240109130712-11603560817a
+	github.com/stretchr/testify v1.8.4
+)
 
 require (
 	github.com/NVIDIA/go-nvml v0.12.0-1.0.20231020145430-e06766c5e74f // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/uuid v1.4.0 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
 replace (

go.sum

@@ -16,6 +16,7 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

gpuallocator/device.go

@@ -78,7 +78,7 @@ func NewDevices(opts ...Option) (DeviceList, error) {
 		opt(o)
 	}
 	if o.nvmllib == nil {
-		o.nvmllib = nvml.New()
+		o.nvmllib = nvmlNew()
 	}
 	if o.devicelib == nil {
 		o.devicelib = device.New(
@@ -139,6 +139,9 @@ func (o *deviceListBuilder) build() (DeviceList, error) {
 
 // NewDevicesFrom creates a list of Devices from the specific set of GPU uuids passed in.
 func NewDevicesFrom(uuids []string) (DeviceList, error) {
+	if len(uuids) == 0 {
+		return DeviceList{}, nil
+	}
 	devices, err := NewDevices()
 	if err != nil {
 		return nil, err
@@ -147,14 +150,9 @@ func NewDevicesFrom(uuids []string) (DeviceList, error) {
 }
 
 // Filter filters out the selected devices from the list.
-// If the supplied list of uuids is nil, no filtering is performed.
 // Note that the specified uuids must exist in the list of devices.
 func (d DeviceList) Filter(uuids []string) (DeviceList, error) {
-	if uuids == nil {
-		return d, nil
-	}
-
-	filtered := []*Device{}
+	var filtered DeviceList
 	for _, uuid := range uuids {
 		for _, device := range d {
 			if device.UUID == uuid {
@@ -254,3 +252,8 @@ func (ds DeviceSet) SortedSlice() []*Device {
 
 	return devices
 }
+
+// nvmlNew is implemented as a function here to allow for this to be replaced for testing.
+var nvmlNew = func() nvml.Interface {
+	return nvml.New()
+}

gpuallocator/device_test.go

@@ -0,0 +1,92 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package gpuallocator
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/go-nvlib/pkg/nvml"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDeviceListFilter(t *testing.T) {
+	singleDeviceNVML := &nvml.InterfaceMock{
+		InitFunc: func() nvml.Return {
+			return nvml.SUCCESS
+		},
+		ShutdownFunc: func() nvml.Return {
+			return nvml.SUCCESS
+		},
+		DeviceGetCountFunc: func() (int, nvml.Return) {
+			return 1, nvml.SUCCESS
+		},
+		DeviceGetHandleByIndexFunc: func(Index int) (nvml.Device, nvml.Return) {
+			device := &nvml.DeviceMock{
+				GetNameFunc: func() (string, nvml.Return) {
+					return "Device0", nvml.SUCCESS
+				},
+				GetUUIDFunc: func() (string, nvml.Return) {
+					return "GPU-0", nvml.SUCCESS
+				},
+				GetPciInfoFunc: func() (nvml.PciInfo, nvml.Return) {
+					return nvml.PciInfo{}, nvml.SUCCESS
+				},
+			}
+			return device, nvml.SUCCESS
+		},
+	}
+
+	testCases := []struct {
+		description        string
+		uuids              []string
+		nvmllib            nvml.Interface
+		expectedDeviceList DeviceList
+		expectedError      error
+	}{
+		{
+			description:        "nil uuids returns empty list",
+			nvmllib:            singleDeviceNVML,
+			expectedDeviceList: DeviceList{},
+		},
+		{
+			description:        "empty uuids returns empty list",
+			uuids:              []string{},
+			nvmllib:            singleDeviceNVML,
+			expectedDeviceList: DeviceList{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			defer setNVMLNewDuringTest(tc.nvmllib)()
+			deviceList, err := NewDevicesFrom(tc.uuids)
+			require.ErrorIs(t, tc.expectedError, err)
+			require.EqualValues(t, tc.expectedDeviceList, deviceList)
+		})
+	}
+}
+
+func setNVMLNewDuringTest(to nvml.Interface) func() {
+	original := nvmlNew
+	nvmlNew = func() nvml.Interface {
+		return to
+	}
+
+	return func() {
+		nvmlNew = original
+	}
+}