From 967344f1952efb61e3159464a734d6cffbcd37ed Mon Sep 17 00:00:00 2001
From: pdrum <pedram.esmaeeli@gmail.com>
Date: Thu, 16 Aug 2018 19:43:17 +0430
Subject: [PATCH] bug: Prevent wiping down output file on error.

---
 main.go               | 28 ++++++++++----------------
 out/lazy_file.go      | 38 ++++++++++++++++++++++++++++++++++
 out/lazy_file_test.go | 47 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 96 insertions(+), 17 deletions(-)
 create mode 100644 out/lazy_file.go
 create mode 100644 out/lazy_file_test.go

diff --git a/main.go b/main.go
index 067a3d8..fe06a6c 100644
--- a/main.go
+++ b/main.go
@@ -8,9 +8,9 @@ import (
 	"io/ioutil"
 	"net/http"
 	"os"
-	"path"
 	"strings"
 
+	"github.com/cheekybits/genny/out"
 	"github.com/cheekybits/genny/parse"
 )
 
@@ -61,22 +61,7 @@ func main() {
 		fatal(exitcodeInvalidTypeSet, err)
 	}
 
-	var outWriter io.Writer
-	if len(*out) > 0 {
-		err := os.MkdirAll(path.Dir(*out), 0755)
-		if err != nil {
-			fatal(exitcodeDestFileFailed, err)
-		}
-
-		outFile, err := os.Create(*out)
-		if err != nil {
-			fatal(exitcodeDestFileFailed, err)
-		}
-		defer outFile.Close()
-		outWriter = outFile
-	} else {
-		outWriter = os.Stdout
-	}
+	outWriter := newWriter(*out)
 
 	if strings.ToLower(args[0]) == "get" {
 		if len(args) != 3 {
@@ -139,6 +124,15 @@ Flags:`)
 	flag.PrintDefaults()
 }
 
+func newWriter(fileName string) io.Writer {
+	if fileName == "" {
+		return os.Stdout
+	}
+	lf := &out.LazyFile{FileName: fileName}
+	defer lf.Close()
+	return lf
+}
+
 func fatal(code int, a ...interface{}) {
 	fmt.Println(a...)
 	os.Exit(code)
diff --git a/out/lazy_file.go b/out/lazy_file.go
new file mode 100644
index 0000000..8591150
--- /dev/null
+++ b/out/lazy_file.go
@@ -0,0 +1,38 @@
+package out
+
+import (
+	"os"
+	"path"
+)
+
+// LazyFile is an io.WriteCloser which defers creation of the file it is supposed to write in
+// till the first call to its write function in order to prevent creation of file, if no write
+// is supposed to happen.
+type LazyFile struct {
+	FileName string
+	written  []byte
+	file     *os.File
+}
+
+// Close closes the file if it is created. Returns nil if no file is created.
+func (lw *LazyFile) Close() error {
+	if lw.file != nil {
+		return lw.file.Close()
+	}
+	return nil
+}
+
+// Write writes to the specified file and creates the file first time it is called.
+func (lw *LazyFile) Write(p []byte) (int, error) {
+	if lw.file == nil {
+		err := os.MkdirAll(path.Dir(lw.FileName), 0755)
+		if err != nil {
+			return 0, err
+		}
+		lw.file, err = os.Create(lw.FileName)
+		if err != nil {
+			return 0, err
+		}
+	}
+	return lw.file.Write(p)
+}
diff --git a/out/lazy_file_test.go b/out/lazy_file_test.go
new file mode 100644
index 0000000..cfca49d
--- /dev/null
+++ b/out/lazy_file_test.go
@@ -0,0 +1,47 @@
+package out_test
+
+import (
+	"github.com/cheekybits/genny/out"
+	"github.com/stretchr/testify/assert"
+	"io/ioutil"
+	"os"
+	"testing"
+)
+
+const testFileName = "test-file.go"
+
+func tearDown() {
+	var err = os.Remove(testFileName)
+	if err != nil && !os.IsNotExist(err) {
+		panic("Could not delete test file")
+	}
+}
+
+func assertFileContains(t *testing.T, expected string) {
+	file, err := os.Open(testFileName)
+	if err != nil {
+		panic(err)
+	}
+	fileBytes, err := ioutil.ReadAll(file)
+	if err != nil {
+		panic(err)
+	}
+	assert.Equal(t, expected, string(fileBytes), "File contents not written properly")
+}
+
+func TestMultipleWrites(t *testing.T) {
+	defer tearDown()
+	lf := out.LazyFile{FileName: testFileName}
+	defer lf.Close()
+	lf.Write([]byte("Word1"))
+	lf.Write([]byte("Word2"))
+	assertFileContains(t, "Word1Word2")
+}
+
+func TestNoWrite(t *testing.T) {
+	defer tearDown()
+	lf := out.LazyFile{FileName: testFileName}
+	defer lf.Close()
+	_, err := os.Stat(testFileName)
+	assert.True(t, os.IsNotExist(err), "Expected file not to be created")
+}