From 14df66d915d2ba2ef69b9cd59dcca0584be19229 Mon Sep 17 00:00:00 2001
From: Peter Stace <peterstace@gmail.com>
Date: Tue, 14 Jul 2015 20:58:58 +1000
Subject: [PATCH 1/2] Format generated source code before writing it to the
 destination

---
 mockgen/mockgen.go | 25 ++++++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index 6371f7da..5226c3b3 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -19,8 +19,10 @@ package main
 // TODO: This does not support embedding package-local interfaces in a separate file.
 
 import (
+	"bytes"
 	"flag"
 	"fmt"
+	"go/format"
 	"go/token"
 	"io"
 	"log"
@@ -86,9 +88,7 @@ func main() {
 		packageName = "mock_" + sanitize(pkg.Name)
 	}
 
-	g := generator{
-		w: dst,
-	}
+	g := generator{}
 	if *source != "" {
 		g.filename = *source
 	} else {
@@ -98,6 +98,9 @@ func main() {
 	if err := g.Generate(pkg, packageName); err != nil {
 		log.Fatalf("Failed generating mock: %v", err)
 	}
+	if _, err := dst.Write(g.Format()); err != nil {
+		log.Fatalf("Failed writing to destination: %v", err)
+	}
 }
 
 func usage() {
@@ -123,7 +126,7 @@ Example:
 `
 
 type generator struct {
-	w      io.Writer
+	buf    bytes.Buffer
 	indent string
 
 	filename                  string // may be empty
@@ -133,7 +136,7 @@ type generator struct {
 }
 
 func (g *generator) p(format string, args ...interface{}) {
-	fmt.Fprintf(g.w, g.indent+format+"\n", args...)
+	fmt.Fprintf(&g.buf, g.indent+format+"\n", args...)
 }
 
 func (g *generator) in() {
@@ -411,3 +414,15 @@ func (g *generator) GenerateMockRecorderMethod(mockType string, m *model.Method)
 	g.p("}")
 	return nil
 }
+
+// Format gives a go-formatted copy of the buffer.
+func (g *generator) Format() []byte {
+	src, err := format.Source(g.buf.Bytes())
+	if err != nil {
+		// This should never happen, but if it does, give the user a warning
+		// and return the contents of the unmodified buffer.
+		log.Printf("Warning: generated source could not be go-formatted: %s", err)
+		return g.buf.Bytes()
+	}
+	return src
+}

From 34906d101d1998ff776654dba203ad02bf204117 Mon Sep 17 00:00:00 2001
From: Peter Stace <peterstace@gmail.com>
Date: Wed, 15 Jul 2015 20:47:27 +1000
Subject: [PATCH 2/2] Log fatally if the generated source cannot be formatted
 in the standard Go style.

---
 mockgen/mockgen.go | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/mockgen/mockgen.go b/mockgen/mockgen.go
index 5226c3b3..6698b2a8 100644
--- a/mockgen/mockgen.go
+++ b/mockgen/mockgen.go
@@ -88,7 +88,7 @@ func main() {
 		packageName = "mock_" + sanitize(pkg.Name)
 	}
 
-	g := generator{}
+	g := new(generator)
 	if *source != "" {
 		g.filename = *source
 	} else {
@@ -98,7 +98,7 @@ func main() {
 	if err := g.Generate(pkg, packageName); err != nil {
 		log.Fatalf("Failed generating mock: %v", err)
 	}
-	if _, err := dst.Write(g.Format()); err != nil {
+	if _, err := dst.Write(g.Output()); err != nil {
 		log.Fatalf("Failed writing to destination: %v", err)
 	}
 }
@@ -415,14 +415,11 @@ func (g *generator) GenerateMockRecorderMethod(mockType string, m *model.Method)
 	return nil
 }
 
-// Format gives a go-formatted copy of the buffer.
-func (g *generator) Format() []byte {
+// Output returns the generator's output, formatted in the standard Go style.
+func (g *generator) Output() []byte {
 	src, err := format.Source(g.buf.Bytes())
 	if err != nil {
-		// This should never happen, but if it does, give the user a warning
-		// and return the contents of the unmodified buffer.
-		log.Printf("Warning: generated source could not be go-formatted: %s", err)
-		return g.buf.Bytes()
+		log.Fatalf("Failed to format generated source code: %s\n%s", err, g.buf.String())
 	}
 	return src
 }