Remove duplicate graph writing code

1 files changed, 16 insertions, 46 deletions

diff --git a/cli/main.go b/cli/main.go
index 33fb7d7..6e12848 100644
--- a/cli/main.go
+++ b/cli/main.go
@@ -3,7 +3,7 @@ package main
 import (
 	"flag"
 	"fmt"
-	"log"
+	"errors"
 	"os"
 	"os/exec"
 	"strings"
@@ -18,53 +18,33 @@ func openBrowser(urlString string) {
 }
 
 func logState(oldState string, newState string, data string) {
-	log.Printf("State: %s -> State: %s with data %s\n", oldState, newState, data)
+	fmt.Printf("State: %s -> State: %s with data %s\n", oldState, newState, data)
 
 	if newState == "SERVER_OAUTH_STARTED" {
 		openBrowser(data)
 	}
 }
 
-func getGraphviz(fsm *eduvpn.FSM, graph string) string {
-	if fsm == nil {
-		return graph
-	}
+func writeGraph(filename string) error {
+	state := eduvpn.GetVPNState()
 
-	for name, state := range fsm.States {
-		for _, transition := range state.Transition {
-			graph += "\n" + "cluster_" + name.String() + "-> cluster_" + transition.String()
-		}
+	state.InitializeFSM()
 
-		graph += "\nsubgraph cluster_" + name.String() + "{\n"
-		if (state.Locked) {
-			graph += "bgcolor=\"red\"\n"
-		}
-		if (fsm.Current == name) {
-			graph += "style=\"bold\"\n"
-			graph += "color=\"blue\"\n"
-		} else {
-			graph += "style=\"\"\n"
-			graph += "color=\"\"\n"
-		}
-		graph += "label=" + name.String()
-		graph = getGraphviz(state.Sub, graph)
-		graph += "\n}"
-	}
-	return graph
-}
+	graph := state.GenerateGraph()
 
-func generateGraph() string {
-	state := eduvpn.GetVPNState()
+	f, err := os.Create(filename)
 
-	state.InitializeFSM()
+	if err != nil {
+		return errors.New(fmt.Sprintf("Failed to create file %s with error %v", filename, err))
+	}
 
+	defer f.Close()
 
-	graph := "digraph fsm {\n"
-	graph += "nodesep=2"
-	graph = getGraphviz(state.FSM, graph)
-	graph += "\n}"
+	f.WriteString(graph)
 
-	return graph
+	fmt.Printf("Graph written to file: %s, use 'fdp %s -Tsvg > graph.svg' from graphviz to save to a svg file called graph.svg\n", filename, filename)
+
+	return nil
 }
 
 func main() {
@@ -74,17 +54,7 @@ func main() {
 
 	fileGraphString := *fileGraph
 	if fileGraphString != "" {
-		f, err := os.Create(fileGraphString)
-
-		if err != nil {
-			log.Fatalf("Failed to create file %s with error %v", fileGraphString, err)
-		}
-
-		defer f.Close()
-
-		f.WriteString(generateGraph())
-
-		log.Printf("Graph written to file: %s, use 'fdp %s -Tsvg > graph.svg' from graphviz to save to a svg file called graph.svg\n", fileGraphString, fileGraphString)
+		writeGraph(fileGraphString)
 		return
 	}
 	urlString := *urlArg