From efb9e6987c00c280f4c3c624a105a9f6c27f2122 Mon Sep 17 00:00:00 2001
From: Kevin S Kirkup <kevin.kirkup@pureport.com>
Date: Tue, 15 Dec 2020 18:25:27 -0500
Subject: [PATCH] Add support for using template to generate tests (#752)

* Add support for using template to generate tests

Fixes #749

* Add unit test

* Fix test - package is foo_bar

* Fix integration test

For Test templates, the import path for the package being tests was
incorrect. Updating to check for a string that ends with the package
name.
---
 ginkgo/generate_command.go     | 25 ++++++++++++++++++-------
 integration/subcommand_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/ginkgo/generate_command.go b/ginkgo/generate_command.go
index f4318d948..288df7797 100644
--- a/ginkgo/generate_command.go
+++ b/ginkgo/generate_command.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"flag"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -12,11 +13,15 @@ import (
 )
 
 func BuildGenerateCommand() *Command {
-	var agouti, noDot, internal bool
+	var (
+		agouti, noDot, internal bool
+		customTestFile          string
+	)
 	flagSet := flag.NewFlagSet("generate", flag.ExitOnError)
 	flagSet.BoolVar(&agouti, "agouti", false, "If set, generate will generate a test file for writing Agouti tests")
 	flagSet.BoolVar(&noDot, "nodot", false, "If set, generate will generate a test file that does not . import ginkgo and gomega")
 	flagSet.BoolVar(&internal, "internal", false, "If set, generate will generate a test file that uses the regular package name")
+	flagSet.StringVar(&customTestFile, "template", "", "If specified, generate will use the contents of the file passed as the test file template")
 
 	return &Command{
 		Name:         "generate",
@@ -28,7 +33,7 @@ func BuildGenerateCommand() *Command {
 			"Accepts the following flags:",
 		},
 		Command: func(args []string, additionalArgs []string) {
-			generateSpec(args, agouti, noDot, internal)
+			generateSpec(args, agouti, noDot, internal, customTestFile)
 		},
 	}
 }
@@ -81,9 +86,9 @@ type specData struct {
 	ImportPackage     bool
 }
 
-func generateSpec(args []string, agouti, noDot, internal bool) {
+func generateSpec(args []string, agouti, noDot, internal bool, customTestFile string) {
 	if len(args) == 0 {
-		err := generateSpecForSubject("", agouti, noDot, internal)
+		err := generateSpecForSubject("", agouti, noDot, internal, customTestFile)
 		if err != nil {
 			fmt.Println(err.Error())
 			fmt.Println("")
@@ -95,7 +100,7 @@ func generateSpec(args []string, agouti, noDot, internal bool) {
 
 	var failed bool
 	for _, arg := range args {
-		err := generateSpecForSubject(arg, agouti, noDot, internal)
+		err := generateSpecForSubject(arg, agouti, noDot, internal, customTestFile)
 		if err != nil {
 			failed = true
 			fmt.Println(err.Error())
@@ -107,7 +112,7 @@ func generateSpec(args []string, agouti, noDot, internal bool) {
 	}
 }
 
-func generateSpecForSubject(subject string, agouti, noDot, internal bool) error {
+func generateSpecForSubject(subject string, agouti, noDot, internal bool, customTestFile string) error {
 	packageName, specFilePrefix, formattedName := getPackageAndFormattedName()
 	if subject != "" {
 		specFilePrefix = formatSubject(subject)
@@ -140,7 +145,13 @@ func generateSpecForSubject(subject string, agouti, noDot, internal bool) error
 	defer f.Close()
 
 	var templateText string
-	if agouti {
+	if customTestFile != "" {
+		tpl, err := ioutil.ReadFile(customTestFile)
+		if err != nil {
+			panic(err.Error())
+		}
+		templateText = string(tpl)
+	} else if agouti {
 		templateText = agoutiSpecText
 	} else {
 		templateText = specText
diff --git a/integration/subcommand_test.go b/integration/subcommand_test.go
index f46777b69..11ba8dbb3 100644
--- a/integration/subcommand_test.go
+++ b/integration/subcommand_test.go
@@ -175,6 +175,36 @@ var _ = Describe("Subcommand", func() {
 			})
 		})
 
+		Context("with template argument", func() {
+			It("should generate a test file using a template", func() {
+				templateFile := filepath.Join(pkgPath, ".generate")
+				ioutil.WriteFile(templateFile, []byte(`package {{.Package}}
+				import (
+					{{if .IncludeImports}}. "github.com/onsi/ginkgo"{{end}}
+					{{if .IncludeImports}}. "github.com/onsi/gomega"{{end}}
+				
+					{{if .ImportPackage}}"{{.PackageImportPath}}"{{end}}
+				)
+				
+				var _ = Describe("{{.Subject}}", func() {
+					// This is a {{.Package}} test
+				})`), 0666)
+				session := startGinkgo(pkgPath, "generate", "--template", ".generate")
+				Eventually(session).Should(gexec.Exit(0))
+				output := session.Out.Contents()
+
+				Ω(output).Should(ContainSubstring("foo_bar_test.go"))
+
+				content, err := ioutil.ReadFile(filepath.Join(pkgPath, "foo_bar_test.go"))
+				Ω(err).ShouldNot(HaveOccurred())
+				Ω(content).Should(ContainSubstring("package foo_bar_test"))
+				Ω(content).Should(ContainSubstring(`. "github.com/onsi/ginkgo"`))
+				Ω(content).Should(ContainSubstring(`. "github.com/onsi/gomega"`))
+				Ω(content).Should(ContainSubstring(`/foo_bar"`))
+				Ω(content).Should(ContainSubstring("// This is a foo_bar_test test"))
+			})
+		})
+
 		Context("with an argument of the form: foo", func() {
 			It("should generate a test file named after the argument", func() {
 				session := startGinkgo(pkgPath, "generate", "baz_buzz")