From a63125b54ffb3ef6a10c5eb27f9fd1096838dee3 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 14 Jul 2014 09:00:01 -0700
Subject: [PATCH] digraph: more idiomatic writedot api

---
 digraph/graphviz.go      | 10 +++++--
 digraph/graphviz_test.go | 57 ++++++++++++++++------------------------
 2 files changed, 30 insertions(+), 37 deletions(-)

diff --git a/digraph/graphviz.go b/digraph/graphviz.go
index 2908dfcd8579..cf0c4912f7c5 100644
--- a/digraph/graphviz.go
+++ b/digraph/graphviz.go
@@ -7,11 +7,15 @@ import (
 
 // GenerateDot is used to emit a GraphViz compatible definition
 // for a directed graph. It can be used to dump a .dot file.
-func GenerateDot(nodes []Node, w io.Writer) {
+func WriteDot(w io.Writer, nodes []Node) error {
 	w.Write([]byte("digraph {\n"))
 	defer w.Write([]byte("}\n"))
+
 	for _, n := range nodes {
-		w.Write([]byte(fmt.Sprintf("\t\"%s\";\n", n)))
+		nodeLine := fmt.Sprintf("\t\"%s\";\n", n)
+
+		w.Write([]byte(nodeLine))
+
 		for _, edge := range n.Edges() {
 			target := edge.Tail()
 			line := fmt.Sprintf("\t\"%s\" -> \"%s\" [label=\"%s\"];\n",
@@ -19,4 +23,6 @@ func GenerateDot(nodes []Node, w io.Writer) {
 			w.Write([]byte(line))
 		}
 	}
+
+	return nil
 }
diff --git a/digraph/graphviz_test.go b/digraph/graphviz_test.go
index 17bcb6c7e8a7..fce1ebb6f503 100644
--- a/digraph/graphviz_test.go
+++ b/digraph/graphviz_test.go
@@ -6,7 +6,7 @@ import (
 	"testing"
 )
 
-func Test_GenerateDot(t *testing.T) {
+func TestWriteDot(t *testing.T) {
 	nodes := ParseBasic(`a -> b ; foo
 a -> c
 b -> d
@@ -18,40 +18,27 @@ b -> e
 	}
 
 	buf := bytes.NewBuffer(nil)
-	GenerateDot(nlist, buf)
-
-	out := string(buf.Bytes())
-	if !strings.HasPrefix(out, "digraph {\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.HasSuffix(out, "\n}\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"a\";\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"b\";\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"c\";\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"d\";\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"e\";\n") {
-		t.Fatalf("bad: %v", out)
+	if err := WriteDot(buf, nlist); err != nil {
+		t.Fatalf("err: %s", err)
 	}
-	if !strings.Contains(out, "\n\t\"a\" -> \"b\" [label=\"foo\"];\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"a\" -> \"c\" [label=\"Edge\"];\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"b\" -> \"d\" [label=\"Edge\"];\n") {
-		t.Fatalf("bad: %v", out)
-	}
-	if !strings.Contains(out, "\n\t\"b\" -> \"e\" [label=\"Edge\"];\n") {
-		t.Fatalf("bad: %v", out)
+
+	actual := strings.TrimSpace(string(buf.Bytes()))
+	expected := strings.TrimSpace(writeDotStr)
+	if actual != expected {
+		t.Fatalf("bad: %s", actual)
 	}
 }
+
+const writeDotStr = `
+digraph {
+	"a";
+	"a" -> "b" [label="foo"];
+	"a" -> "c" [label="Edge"];
+	"b";
+	"b" -> "d" [label="Edge"];
+	"b" -> "e" [label="Edge"];
+	"c";
+	"d";
+	"e";
+}
+`