markdown.go

// Copyright (C) 2019 rameshvk. All rights reserved.
// Use of this source code is governed by a MIT-style license
// that can be found in the LICENSE file.

package test

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"io/ioutil"
	"log"
	"runtime"
	"strings"
	"text/template"

	"github.com/russross/blackfriday/v2"
	"golang.org/x/tools/imports"
)

// Markdown generates test code by compiling together snippets in the
// provided markdown file names.  The output is written to the provided
// output file name.
//
// Fenced blocks can also specify the test/example name in the info string.
// The first word of the infostring must be go or golang
//
// A fenced block can "import" other modules by simply placing a comment
// of the form `// import path`.
//
// See https://github.com/tvastar/test/blob/master/testdata/markdown.md
// for a sample markdown and
// https://github.com/tvastar/test/blob/master/testdata/markdown_test.go
// for example generated tests.
func Markdown(src []string, dest, pkg string) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = r.(error)
		}
	}()

	pc := []uintptr{0}
	runtime.Callers(2, pc)
	frame, _ := runtime.CallersFrames(pc).Next()

	var output bytes.Buffer

	im := [][2]string{{"", "testing"}, {"", "fmt"}, {"", "log"}}
	i := info{Package: pkg, Imports: im, Generator: sanitize(frame.File)}

	input := readSources(src)

	count := 0
	opts := blackfriday.WithExtensions(blackfriday.CommonExtensions)
	blackfriday.New(opts).Parse(input).Walk(
		func(n *blackfriday.Node, entering bool) blackfriday.WalkStatus {
			if entering && n.Type == blackfriday.CodeBlock {
				count++
				fence := string(n.Literal)
				title := string(n.CodeBlockData.Info)
				im := formatFence(&output, fence, title, pkg, count)
				i.Imports = append(i.Imports, im...)
			}
			return blackfriday.GoToNext
		})

	var hdr bytes.Buffer
	must(headerTpl.Execute(&hdr, i), "")

	contents := append(hdr.Bytes(), output.Bytes()...)
	p, err := format.Source(contents)
	must(err, string(contents))

	result, err := imports.Process(src[0], p, nil)
	must(err, string(p))

	return ioutil.WriteFile(dest, result, 0644)
}

func readSources(sources []string) []byte {
	var result []byte
	for _, src := range sources {
		input, err := ioutil.ReadFile(src)
		must(err, src)
		result = append(result, input...)
	}
	return result
}

func sanitize(fileName string) string {
	if x := strings.Index(fileName, "github.com"); x >= 0 {
		return fileName[x:]
	}
	return fileName
}

func formatFence(w io.Writer, fence, title, pkg string, count int) [][2]string {
	title = strings.TrimSpace(title)
	if !strings.HasPrefix(title, "go") {
		return nil
	}

	parts := strings.Fields(title)
	name := ""
	if title != "" && len(parts) > 1 {
		name = parts[1]
	}

	if strings.Contains(name, ".") {
		parts = strings.SplitN(name, ".", 2)
		if strings.TrimSpace(parts[0]) != pkg {
			return nil
		}
		name = parts[1]
	}

	if name == "" {
		name = fmt.Sprintf("Example_%d", count)
	}

	if strings.HasPrefix(name, "Test") {
		ioMust(io.WriteString(w, "func "))
		ioMust(io.WriteString(w, name))
		ioMust(io.WriteString(w, "(t *testing.T) {\n"))
		ioMust(io.WriteString(w, fence))
		ioMust(io.WriteString(w, "\n}\n\n"))
	} else if name == "skip" {
	} else if name == "global" {
		ioMust(io.WriteString(w, fence))
	} else {
		ioMust(io.WriteString(w, "func "))
		ioMust(io.WriteString(w, name))
		ioMust(io.WriteString(w, "() {\n"))
		ioMust(io.WriteString(w, fence))
		ioMust(io.WriteString(w, "\n}\n\n"))
	}

	imports := [][2]string{}
	for _, l := range strings.Split(fence, "\n") {
		l = strings.TrimSpace(l)
		if !strings.HasPrefix(l, "//") {
			continue
		}
		l = strings.TrimSpace(l[2:])
		if !strings.HasPrefix(l, "import") {
			continue
		}
		l = strings.TrimSpace(l[len("import"):])
		imports = append(imports, [2]string{"", l})
	}
	return imports
}

func ioMust(n int, err error) {
	if err != nil {
		panic(err)
	}
}

func must(err error, s string) {
	if err != nil {
		log.Println(s)
		panic(err)
	}
}

var headerTpl = template.Must(template.New("code").Parse(`
// Code generated by {{.Generator}}. DO NOT EDIT.

package {{.Package}}

import (
  {{range $import := .Imports}}{{index $import 0}} "{{index $import 1}}"
  {{end -}}
)

`))

type info struct {
	Generator string
	Package   string
	Imports   [][2]string
}