Add test for GHToc.Print

cmd/gh-md-toc/main.go

@@ -66,7 +66,7 @@ func main() {
 		toc := <-ch
 		// #14, check if there's really TOC?
 		if toc != nil {
-			toc.Print()
+			toc.Print(os.Stdout)
 		}
 	}
 
@@ -80,7 +80,8 @@ func main() {
 		defer os.Remove(file.Name())
 
 		check(ioutil.WriteFile(file.Name(), bytes, 0644))
-		ghtoc.NewGHDoc(file.Name(), false, *startDepth, *depth, !*noEscape, *token, *indent, *debug).GetToc().Print()
+		ghtoc.NewGHDoc(file.Name(), false, *startDepth, *depth, !*noEscape, *token, *indent, *debug).
+			GetToc().Print(os.Stdout)
 	}
 
 	if !*hideFooter {

ghdoc.go

@@ -2,6 +2,7 @@ package ghtoc
 
 import (
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"net/url"
@@ -15,11 +16,11 @@ import (
 type GHToc []string
 
 // Print TOC to the console
-func (toc *GHToc) Print() {
+func (toc *GHToc) Print(w io.Writer) {
 	for _, tocItem := range *toc {
-		fmt.Println(tocItem)
+		fmt.Fprint(w, tocItem)
 	}
-	fmt.Println()
+	fmt.Fprintln(w)
 }
 
 // GHDoc GitHub document

ghdoc_test.go

@@ -1,6 +1,9 @@
 package ghtoc
 
-import "testing"
+import (
+	"bytes"
+	"testing"
+)
 
 func Test_IsUrl(t *testing.T) {
 	doc1 := &GHDoc{
@@ -346,3 +349,14 @@ func Test_MinHeaderNumber(t *testing.T) {
 		t.Error("Res :", toc, "\nExpected     :", tocExpected)
 	}
 }
+
+func TestGHToc_Print(t *testing.T) {
+	toc := GHToc{"one", "two"}
+	want := "onetwo\n"
+	var got bytes.Buffer
+	toc.Print(&got)
+
+	if got.String() != want {
+		t.Error("\nGot :", got.String(), "\nWant:", want)
+	}
+}