From 550588665505b7ec980c145b217d87585ffb46ae Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Tue, 15 Aug 2023 14:29:22 +0200
Subject: [PATCH] Use options for NewLibraryLocator

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/discover/graphics.go   |  8 ++++----
 internal/lookup/cuda/cuda.go    |  9 +++------
 internal/lookup/library.go      | 25 +++++++++++++------------
 internal/lookup/library_test.go | 12 ++++++++----
 4 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/internal/discover/graphics.go b/internal/discover/graphics.go
index 73145c89..aeaabe69 100644
--- a/internal/discover/graphics.go
+++ b/internal/discover/graphics.go
@@ -49,10 +49,10 @@ func NewDRMNodesDiscoverer(logger logger.Interface, devices image.VisibleDevices
 
 // NewGraphicsMountsDiscoverer creates a discoverer for the mounts required by graphics tools such as vulkan.
 func NewGraphicsMountsDiscoverer(logger logger.Interface, driverRoot string, nvidiaCTKPath string) (Discover, error) {
-	locator, err := lookup.NewLibraryLocator(logger, driverRoot)
-	if err != nil {
-		return nil, fmt.Errorf("failed to construct library locator: %v", err)
-	}
+	locator := lookup.NewLibraryLocator(
+		lookup.WithLogger(logger),
+		lookup.WithRoot(driverRoot),
+	)
 	libraries := NewMounts(
 		logger,
 		locator,
diff --git a/internal/lookup/cuda/cuda.go b/internal/lookup/cuda/cuda.go
index b95e81b4..0a8620bc 100644
--- a/internal/lookup/cuda/cuda.go
+++ b/internal/lookup/cuda/cuda.go
@@ -58,13 +58,10 @@ func New(opts ...Options) lookup.Locator {
 		c.driverRoot = "/"
 	}
 
-	// TODO: Do we want to set the Count to 1 here?
-	l, _ := lookup.NewLibraryLocator(
-		c.logger,
-		c.driverRoot,
+	c.Locator = lookup.NewLibraryLocator(
+		lookup.WithLogger(c.logger),
+		lookup.WithRoot(c.driverRoot),
 	)
-
-	c.Locator = l
 	return c
 }
 
diff --git a/internal/lookup/library.go b/internal/lookup/library.go
index 73ad6984..4a6bb722 100644
--- a/internal/lookup/library.go
+++ b/internal/lookup/library.go
@@ -30,12 +30,9 @@ type ldcacheLocator struct {
 
 var _ Locator = (*ldcacheLocator)(nil)
 
-// NewLibraryLocator creates a library locator using the specified logger.
-func NewLibraryLocator(logger logger.Interface, root string) (Locator, error) {
-	// We construct a symlink locator for expected library locations.
-	symlinkLocator := NewSymlinkLocator(
-		WithLogger(logger),
-		WithRoot(root),
+// NewLibraryLocator creates a library locator using the specified options.
+func NewLibraryLocator(opts ...Option) Locator {
+	opts = append(opts,
 		WithSearchPaths([]string{
 			"/",
 			"/usr/lib64",
@@ -50,24 +47,28 @@ func NewLibraryLocator(logger logger.Interface, root string) (Locator, error) {
 			"/lib/aarch64-linux-gnu/nvidia/current",
 		}...),
 	)
+	// We construct a symlink locator for expected library locations.
+	symlinkLocator := NewSymlinkLocator(opts...)
 
 	l := First(
 		symlinkLocator,
-		newLdcacheLocator(logger, root),
+		newLdcacheLocator(opts...),
 	)
-	return l, nil
+	return l
 }
 
-func newLdcacheLocator(logger logger.Interface, root string) Locator {
-	cache, err := ldcache.New(logger, root)
+func newLdcacheLocator(opts ...Option) Locator {
+	b := newBuilder(opts...)
+
+	cache, err := ldcache.New(b.logger, b.root)
 	if err != nil {
 		// If we failed to open the LDCache, we default to a symlink locator.
-		logger.Warningf("Failed to load ldcache: %v", err)
+		b.logger.Warningf("Failed to load ldcache: %v", err)
 		return nil
 	}
 
 	return &ldcacheLocator{
-		logger: logger,
+		logger: b.logger,
 		cache:  cache,
 	}
 }
diff --git a/internal/lookup/library_test.go b/internal/lookup/library_test.go
index 5682281d..21485dca 100644
--- a/internal/lookup/library_test.go
+++ b/internal/lookup/library_test.go
@@ -43,7 +43,10 @@ func TestLDCacheLocator(t *testing.T) {
 	require.NoError(t, os.Symlink(versionLib, sonameLink))
 	require.NoError(t, os.Symlink(sonameLink, soLink))
 
-	lut := newLdcacheLocator(logger, testDir)
+	lut := newLdcacheLocator(
+		WithLogger(logger),
+		WithRoot(testDir),
+	)
 
 	testCases := []struct {
 		description   string
@@ -94,7 +97,6 @@ func TestLDCacheLocator(t *testing.T) {
 			require.EqualValues(t, tc.expected, cleanedCandidates)
 		})
 	}
-
 }
 
 func TestLibraryLocator(t *testing.T) {
@@ -125,8 +127,10 @@ func TestLibraryLocator(t *testing.T) {
 	require.NoError(t, os.Symlink(libTarget1, source1))
 	require.NoError(t, os.Symlink(source1, source2))
 
-	lut, err := NewLibraryLocator(logger, testDir)
-	require.NoError(t, err)
+	lut := NewLibraryLocator(
+		WithLogger(logger),
+		WithRoot(testDir),
+	)
 
 	testCases := []struct {
 		description   string