From e848aa78134a07d81bbbbfc5ab87d28d71c08471 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Mon, 25 Jul 2022 10:13:54 +0200
Subject: [PATCH] Set toolkit root as flag

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 test/container/toolkit_test.sh     |  4 ++--
 tools/container/toolkit/toolkit.go | 25 ++++++++++++++-----------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/test/container/toolkit_test.sh b/test/container/toolkit_test.sh
index 77df5fb0..84df4cac 100644
--- a/test/container/toolkit_test.sh
+++ b/test/container/toolkit_test.sh
@@ -23,7 +23,7 @@ testing::toolkit::install() {
 		READLINK="greadlink"
 	fi
 
-	testing::docker_run::toolkit::shell 'toolkit install /usr/local/nvidia/toolkit'
+	testing::docker_run::toolkit::shell 'toolkit install --toolkit-root=/usr/local/nvidia/toolkit'
 	docker run --rm -v "${shared_dir}:/work" alpine sh -c "chown -R ${uid}:${gid} /work/"
 
 	# Ensure toolkit dir is correctly setup
@@ -66,7 +66,7 @@ testing::toolkit::install() {
 testing::toolkit::delete() {
 	testing::docker_run::toolkit::shell 'mkdir -p /usr/local/nvidia/delete-toolkit'
 	testing::docker_run::toolkit::shell 'touch /usr/local/nvidia/delete-toolkit/test.file'
-	testing::docker_run::toolkit::shell 'toolkit delete /usr/local/nvidia/delete-toolkit'
+	testing::docker_run::toolkit::shell 'toolkit delete --toolkit-root=/usr/local/nvidia/delete-toolkit'
 
 	test ! -z "$(ls -A "${shared_dir}/usr/local/nvidia")"
 	test ! -e "${shared_dir}/usr/local/nvidia/delete-toolkit"
diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go
index 23d892f3..e8f7fde8 100644
--- a/tools/container/toolkit/toolkit.go
+++ b/tools/container/toolkit/toolkit.go
@@ -63,7 +63,7 @@ func main() {
 	install.Usage = "Install the components of the NVIDIA container toolkit"
 	install.ArgsUsage = "<toolkit_directory>"
 	install.Before = func(c *cli.Context) error {
-		return parseArgs(c, &opts)
+		return validateOptions(c, &opts)
 	}
 	install.Action = func(c *cli.Context) error {
 		return Install(c, &opts)
@@ -75,7 +75,7 @@ func main() {
 	delete.Usage = "Delete the NVIDIA container toolkit"
 	delete.ArgsUsage = "<toolkit_directory>"
 	delete.Before = func(c *cli.Context) error {
-		return parseArgs(c, &opts)
+		return validateOptions(c, &opts)
 	}
 	delete.Action = func(c *cli.Context) error {
 		return Delete(c, &opts)
@@ -111,10 +111,18 @@ func main() {
 			Destination: &opts.ContainerCLIDebug,
 			EnvVars:     []string{"NVIDIA_CONTAINER_CLI_DEBUG"},
 		},
+		&cli.StringFlag{
+			Name:        "toolkit-root",
+			Usage:       "The directory where the NVIDIA Container toolkit is to be installed",
+			Required:    true,
+			Destination: &opts.toolkitRoot,
+			EnvVars:     []string{"TOOLKIT_ROOT"},
+		},
 	}
 
 	// Update the subcommand flags with the common subcommand flags
 	install.Flags = append([]cli.Flag{}, flags...)
+	delete.Flags = append([]cli.Flag{}, flags...)
 
 	// Run the top-level CLI
 	if err := c.Run(os.Args); err != nil {
@@ -122,16 +130,11 @@ func main() {
 	}
 }
 
-// parseArgs parses the command line arguments to the CLI
-func parseArgs(c *cli.Context, opts *options) error {
-	args := c.Args()
-
-	log.Infof("Parsing arguments: %v", args.Slice())
-	if c.NArg() != 1 {
-		return fmt.Errorf("incorrect number of arguments")
+// validateOptions checks whether the specified options are valid
+func validateOptions(c *cli.Context, opts *options) error {
+	if opts.toolkitRoot == "" {
+		return fmt.Errorf("invalid --toolkit-root option: %v", opts.toolkitRoot)
 	}
-	opts.toolkitRoot = args.Get(0)
-	log.Infof("Successfully parsed arguments")
 
 	return nil
 }