diff --git a/cmd/gh-md-toc/main.go b/cmd/gh-md-toc/main.go index 200e206..27e6de8 100644 --- a/cmd/gh-md-toc/main.go +++ b/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 { diff --git a/ghdoc.go b/ghdoc.go index 8a59a2f..2a53ea2 100644 --- a/ghdoc.go +++ b/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 diff --git a/ghdoc_test.go b/ghdoc_test.go index 4846f2a..81d29d1 100644 --- a/ghdoc_test.go +++ b/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) + } +}