From 76018d282e3967cea7300d356708e711939da6f7 Mon Sep 17 00:00:00 2001
From: Christopher Desiniotis <cdesiniotis@nvidia.com>
Date: Fri, 9 Jun 2023 16:06:32 -0700
Subject: [PATCH 1/2] Allow clients of the pciids API to set the pci.ids
 filepath

Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
---
 pkg/pciids/pciids.go | 41 +++++++++++++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 8 deletions(-)

diff --git a/pkg/pciids/pciids.go b/pkg/pciids/pciids.go
index 8af94d9..7a446ab 100644
--- a/pkg/pciids/pciids.go
+++ b/pkg/pciids/pciids.go
@@ -172,24 +172,48 @@ type parser struct {
 	}
 }
 
+// Various locations of pci.ids for different distributions. These may be more
+// up to date then the embedded pci.ids db
+var defaultPCIdbPaths = []string{
+	"/usr/share/misc/pci.ids",   // Ubuntu
+	"/usr/local/share/pci.ids",  // RHEL like with manual update
+	"/usr/share/hwdata/pci.ids", // RHEL like
+	"/usr/share/pci.ids",        // SUSE
+}
+
 // This is a fallback if all of the locations fail
 //go:embed default_pci.ids
 var defaultPCIdb []byte
 
 // NewDB Parse the PCI DB in its default locations or use the default
 // builtin pci.ids db.
-func NewDB() Interface {
-	// Various locations of pci.ids for differente distributions these may be more
-	// up to date then the embedded pci.ids db
-	pcidbs := []string{
-		"/usr/share/misc/pci.ids",   // Ubuntu
-		"/usr/local/share/pci.ids",  // RHEL like with manual update
-		"/usr/share/hwdata/pci.ids", // RHEL like
-		"/usr/share/pci.ids",        // SUSE
+func NewDB(opts ...Option) Interface {
+	db := &pcidb{}
+	for _, opt := range opts {
+		opt(db)
 	}
+
+	pcidbs := defaultPCIdbPaths
+	if db.path != "" {
+		pcidbs = append([]string{db.path}, defaultPCIdbPaths...)
+	}
+
 	return newParser(pcidbs).parse()
 }
 
+// Option defines a function for passing options to the NewDB() call
+type Option func(*pcidb)
+
+// WithFilePath provides an Option to set the file path
+// for the pciids database used by pciids interface.
+// The file path provided takes precedence over all other
+// paths.
+func WithFilePath(path string) Option {
+	return func(db *pcidb) {
+		db.path = path
+	}
+}
+
 // newParser will attempt to read the db pci.ids from well known places or fall
 // back to an internal db
 func newParser(pcidbs []string) *parser {
@@ -247,6 +271,7 @@ func (d *pcidb) GetClassName(classID uint32) string {
 type pcidb struct {
 	vendors map[uint16]vendor
 	classes map[uint32]class
+	path    string
 }
 
 // vendor PCI vendors/devices/subVendors/SubDevices

From 066d8f30bc53e5a4697d2fae727b80c789b6e949 Mon Sep 17 00:00:00 2001
From: Christopher Desiniotis <cdesiniotis@nvidia.com>
Date: Fri, 9 Jun 2023 17:27:31 -0700
Subject: [PATCH 2/2] Allow options to be passed when creating an instance of
 the nvpci interface

Signed-off-by: Christopher Desiniotis <cdesiniotis@nvidia.com>
---
 pkg/nvmdev/nvmdev.go |  4 +++-
 pkg/nvpci/mock.go    |  2 +-
 pkg/nvpci/nvpci.go   | 32 ++++++++++++++++++++++++++------
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/pkg/nvmdev/nvmdev.go b/pkg/nvmdev/nvmdev.go
index fa2c3fc..24ec228 100644
--- a/pkg/nvmdev/nvmdev.go
+++ b/pkg/nvmdev/nvmdev.go
@@ -380,5 +380,7 @@ func (p *ParentDevice) GetAvailableMDEVInstances(mdevType string) (int, error) {
 func newNvidiaPCIDeviceFromPath(devicePath string) (*nvpci.NvidiaPCIDevice, error) {
 	root := filepath.Dir(devicePath)
 	address := filepath.Base(devicePath)
-	return nvpci.NewFrom(root).GetGPUByPciBusID(address)
+	return nvpci.New(
+		nvpci.WithPCIDevicesRoot(root),
+	).GetGPUByPciBusID(address)
 }
diff --git a/pkg/nvpci/mock.go b/pkg/nvpci/mock.go
index 0f3df3a..8f35fb6 100644
--- a/pkg/nvpci/mock.go
+++ b/pkg/nvpci/mock.go
@@ -44,7 +44,7 @@ func NewMockNvpci() (mock *MockNvpci, rerr error) {
 	}()
 
 	mock = &MockNvpci{
-		NewFrom(rootDir).(*nvpci),
+		New(WithPCIDevicesRoot(rootDir)).(*nvpci),
 	}
 
 	return mock, nil
diff --git a/pkg/nvpci/nvpci.go b/pkg/nvpci/nvpci.go
index e6b23b6..1640590 100644
--- a/pkg/nvpci/nvpci.go
+++ b/pkg/nvpci/nvpci.go
@@ -65,6 +65,7 @@ type ResourceInterface interface {
 
 type nvpci struct {
 	pciDevicesRoot string
+	pcidbPath      string
 }
 
 var _ Interface = (*nvpci)(nil)
@@ -124,14 +125,33 @@ func (d *NvidiaPCIDevice) Reset() error {
 }
 
 // New interface that allows us to get a list of all NVIDIA PCI devices
-func New() Interface {
-	return NewFrom(PCIDevicesRoot)
+func New(opts ...Option) Interface {
+	n := &nvpci{}
+	for _, opt := range opts {
+		opt(n)
+	}
+	if n.pciDevicesRoot == "" {
+		n.pciDevicesRoot = PCIDevicesRoot
+	}
+	return n
 }
 
-// NewFrom interface allows us to get a list of all NVIDIA PCI devices at a specific root directory
-func NewFrom(root string) Interface {
-	return &nvpci{
-		pciDevicesRoot: root,
+// Option defines a function for passing options to the New() call
+type Option func(*nvpci)
+
+// WithPCIDevicesRoot provides an Option to set the root path
+// for PCI devices on the system.
+func WithPCIDevicesRoot(root string) Option {
+	return func(n *nvpci) {
+		n.pciDevicesRoot = root
+	}
+}
+
+// WithPCIDatabasePath provides an Option to set the path
+// to the pciids database file.
+func WithPCIDatabasePath(path string) Option {
+	return func(n *nvpci) {
+		n.pcidbPath = path
 	}
 }