main.go

package main

import (
	"os"
	"strings"
	"fmt"
	"log"
	"go/parser"
	"go/token"
	"go/ast"
	"flag"
)

type envVar struct {
	n string
	d *ast.CommentGroup
}

func main() {
	flag.Parse()

	infile := flag.Arg(0)
	outfile := flag.Arg(1)

	if len(outfile) < 4 || outfile[len(outfile)-3:] != ".go" {
		fmt.Fprintf(os.Stderr, "Please specify a valid outfile (second arg)")
		return
	}

	vars, pkg := getEnvConstantsFromFile(infile)

	createLookupFile(vars, pkg, outfile)
	createLookupTestFile(vars, pkg, outfile)
}

func createLookupTestFile(vars []envVar, pkg string, filename string) {
	testFilename := filename[:len(filename)-3]+"_test.go"
	out, err := os.OpenFile(testFilename, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
	if err != nil {
		log.Fatal(err)
	}

	defer out.Close()

	fmt.Fprintln(out,"package",pkg)
	fmt.Fprintln(out,"import \"testing\"")
	fmt.Fprintln(out,"")
	fmt.Fprintln(out, "// DO NOT EDIT")
	fmt.Fprintf(out, "// This file is auto generated by: %s\n", strings.Join(os.Args, " "))
	fmt.Fprintln(out,"")
	fmt.Fprintln(out,"")
	fmt.Fprintln(out, `
func OnPanic(t *testing.T, env string) {
	if err := recover(); err != nil {
		t.Errorf("Missing Env: %v", err)
	}
}`)


	for _, v := range vars {
		fmt.Fprintln(out, createLookupTestFunction(v))
		fmt.Fprintln(out, createMustLookupTestFunction(v))
	}

}


func createLookupFile(vars []envVar, pkg string, filename string) {
	out, err := os.OpenFile(filename, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	fmt.Fprintln(out,"package",pkg)
	fmt.Fprintln(out,"import \"os\"")
	fmt.Fprintln(out,"import \"fmt\"")
	fmt.Fprintln(out,"")
	fmt.Fprintln(out, "// DO NOT EDIT")
	fmt.Fprintf(out, "// This file is auto generated by: %s\n", strings.Join(os.Args, " "))
	fmt.Fprintln(out,"")
	fmt.Fprintln(out,"")
	for _, v := range vars {
		fmt.Fprintln(out, createLookupFunction(v))
		fmt.Fprintln(out, createMustLookupFunction(v))
	}
}

func createCommentBlock(c *ast.CommentGroup) string {
	if c == nil {
		return ""
	}

	var sb strings.Builder

	for _, comment := range c.List {
		sb.WriteString(comment.Text)
		sb.WriteRune('\n')
	}

	return sb.String()
}

func createLookupTestFunction(v envVar) string {
	v.d.Text() 
	funcName := "Get"+v.n[3:] // replace evn with Get
	s := `%sfunc Test%s(t *testing.T) {
	if _, ok := %s(); !ok {
		t.Errorf("Missing Env: %%s", %s)
	}
}
`

	return fmt.Sprintf(s, createCommentBlock(v.d), funcName, funcName ,v.n)
}


func createMustLookupTestFunction(v envVar) string {
	v.d.Text() 
	funcName := "Get"+v.n[3:] // replace evn with Get
	s := `%sfunc TestMust%s(t *testing.T) {
	defer OnPanic(t, %s)
	_ = Must%s()
}
`

	return fmt.Sprintf(s, createCommentBlock(v.d), funcName, v.n, funcName)
}

func createLookupFunction(v envVar) string {
	v.d.Text() 
	funcName := "Get"+v.n[3:] // replace evn with Get
	s := `%sfunc %s() (value string, ok bool) {
	value, ok = os.LookupEnv(%s)
	return
}
`
	return fmt.Sprintf(s, createCommentBlock(v.d), funcName, v.n)
}


func createMustLookupFunction(v envVar) string {
	v.d.Text() 
	funcName := "Get"+v.n[3:] // replace evn with Get
	s := `%sfunc Must%s() (value string) {
	var ok bool
	value, ok = os.LookupEnv(%s)
	if !ok {
		panic(fmt.Errorf("expected environment variable %%s to be set", %s))
	}
	return
}
`

	return fmt.Sprintf(s, createCommentBlock(v.d), funcName, v.n, v.n)
}

func getEnvConstantsFromFile(filepath string) ([]envVar, string) {
	fset := token.NewFileSet()

	f, err := parser.ParseFile(fset, filepath, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}

	var envs []envVar
	var pkg string

	pkg = f.Name.String()
	for _, decl := range f.Decls {
		switch t := decl.(type){
		case *ast.GenDecl:
			if t.Tok == token.CONST {
				for _, spec := range t.Specs {
					switch st := spec.(type) {
					case *ast.ValueSpec:
						name := st.Names[0].String()
						if name[:3] == "env" {
							x := envVar{
								n: name,
								d: st.Doc, 
							}
							envs = append(envs, x)
						}
					}
				}
			}
		}
	}

	return envs, pkg
}