From 618e2465389fa2777b408812efd4aaa5e331724a Mon Sep 17 00:00:00 2001
From: Patrick Hoefler <patrick.hoefler@gmail.com>
Date: Sun, 12 Feb 2023 17:55:38 +0100
Subject: [PATCH] Check if graphviz is installed (#290)

---
 internal/cmd/root.go      | 78 ++++++++++++++++++++++++---------------
 internal/cmd/root_test.go | 12 +++++-
 2 files changed, 60 insertions(+), 30 deletions(-)

diff --git a/internal/cmd/root.go b/internal/cmd/root.go
index 17f60c7..e86c5a4 100644
--- a/internal/cmd/root.go
+++ b/internal/cmd/root.go
@@ -39,7 +39,9 @@ func (d dfgWriter) Write(p []byte) (n int, err error) {
 }
 
 // NewRootCmd creates a new root command.
-func NewRootCmd(dfgWriter io.Writer, inputFS afero.Fs) *cobra.Command {
+func NewRootCmd(
+	dfgWriter io.Writer, inputFS afero.Fs, dotCmd string,
+) *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use:   "dockerfilegraph",
 		Short: "Visualize your multi-stage Dockerfile",
@@ -54,6 +56,13 @@ It creates a visual graph representation of the build process.`,
 				return printVersion(dfgWriter)
 			}
 
+			// Make sure that graphviz is installed.
+			_, err = exec.LookPath(dotCmd)
+			if err != nil {
+				return
+			}
+
+			// Load and parse the Dockerfile.
 			dockerfile, err := dockerfile2dot.LoadAndParseDockerfile(
 				inputFS,
 				filenameFlag,
@@ -90,32 +99,8 @@ It creates a visual graph representation of the build process.`,
 				return
 			}
 
-			var unflattenFile *os.File
 			if unflattenFlag > 0 {
-				unflattenFile, err = os.CreateTemp("", "dockerfile.*.dot")
-				if err != nil {
-					return
-				}
-				defer os.Remove(unflattenFile.Name())
-
-				unflattenCmd := exec.Command(
-					"unflatten",
-					"-l", strconv.FormatUint(uint64(unflattenFlag), 10),
-					"-o", unflattenFile.Name(), dotFile.Name(),
-				)
-				unflattenCmd.Stdout = dfgWriter
-				unflattenCmd.Stderr = dfgWriter
-				err = unflattenCmd.Run()
-				if err != nil {
-					return
-				}
-
-				err = unflattenFile.Close()
-				if err != nil {
-					return
-				}
-
-				err = os.Rename(unflattenFile.Name(), dotFile.Name())
+				err = unflatten(dotFile, dfgWriter)
 				if err != nil {
 					return
 				}
@@ -141,7 +126,7 @@ It creates a visual graph representation of the build process.`,
 			}
 			dotArgs = append(dotArgs, dotFile.Name())
 
-			out, err := exec.Command("dot", dotArgs...).CombinedOutput()
+			out, err := exec.Command(dotCmd, dotArgs...).CombinedOutput()
 			if err != nil {
 				fmt.Fprintf(dfgWriter,
 					`Oh no, something went wrong while generating the graph!
@@ -154,7 +139,7 @@ It creates a visual graph representation of the build process.`,
 					%s`,
 					dotFileContent, string(out),
 				)
-				os.Exit(1)
+				return
 			}
 
 			fmt.Fprintf(dfgWriter, "Successfully created %s\n", filename)
@@ -260,9 +245,44 @@ It creates a visual graph representation of the build process.`,
 	return rootCmd
 }
 
+func unflatten(dotFile *os.File, dfgWriter io.Writer) (err error) {
+	var unflattenFile *os.File
+	unflattenFile, err = os.CreateTemp("", "dockerfile.*.dot")
+	if err != nil {
+		return
+	}
+	defer os.Remove(unflattenFile.Name())
+
+	unflattenCmd := exec.Command(
+		"unflatten",
+		"-l", strconv.FormatUint(uint64(unflattenFlag), 10),
+		"-o", unflattenFile.Name(), dotFile.Name(),
+	)
+	unflattenCmd.Stdout = dfgWriter
+	unflattenCmd.Stderr = dfgWriter
+	err = unflattenCmd.Run()
+	if err != nil {
+		return
+	}
+
+	err = unflattenFile.Close()
+	if err != nil {
+		return
+	}
+
+	err = os.Rename(unflattenFile.Name(), dotFile.Name())
+	if err != nil {
+		return
+	}
+
+	return
+}
+
 // Execute executes the root command.
 func Execute() {
-	err := NewRootCmd(dfgWriter{}, afero.NewOsFs()).Execute()
+	err := NewRootCmd(
+		dfgWriter{}, afero.NewOsFs(), "dot",
+	).Execute()
 	if err != nil {
 		// Cobra prints the error message
 		os.Exit(1)
diff --git a/internal/cmd/root_test.go b/internal/cmd/root_test.go
index cce189a..e85b428 100644
--- a/internal/cmd/root_test.go
+++ b/internal/cmd/root_test.go
@@ -15,6 +15,7 @@ type test struct {
 	name               string
 	cliArgs            []string
 	dockerfileContent  string
+	dotCmd             string
 	wantErr            bool
 	wantOut            string
 	wantOutRegex       string
@@ -95,6 +96,12 @@ It creates a visual graph representation of the build process.
 			wantErr:           true,
 			wantOut:           "Error: file with no instructions\n" + usage + "\n",
 		},
+		{
+			name:    "graphviz not installed",
+			dotCmd:  "dot-not-found-in-path",
+			wantErr: true,
+			wantOut: "Error: exec: \"dot-not-found-in-path\": executable file not found in $PATH\n" + usage + "\n",
+		},
 		{
 			name:    "--max-label-length too small",
 			cliArgs: []string{"--max-label-length", "3"},
@@ -496,7 +503,10 @@ It creates a visual graph representation of the build process.
 		t.Run(tt.name, func(t *testing.T) {
 			buf := new(bytes.Buffer)
 
-			command := cmd.NewRootCmd(buf, inputFS)
+			if tt.dotCmd == "" {
+				tt.dotCmd = "dot"
+			}
+			command := cmd.NewRootCmd(buf, inputFS, tt.dotCmd)
 			command.SetArgs(tt.cliArgs)
 
 			// Redirect Cobra output