func (g Generator) WriteTest(wr io.Writer) error {
var buf bytes.Buffer
if err := testTmpl.Execute(&buf, g); err != nil {
return err
}
return cleanimports.Clean(wr, buf.Bytes())
}
// Write the generated code into the given io.Writer
//
// Returns an error if there is a problem generating the code
func (g *Generator) Write(w io.Writer) error {
buff := bytes.NewBuffer(nil)
// Generate header
if err := headerTpl.Execute(buff, g); err != nil {
return err
}
// Generate funcs
for _, fn := range g.funcs {
if g.Include != nil && !g.Include.MatchString(fn.Name) {
continue
}
if g.Exclude != nil && g.Exclude.MatchString(fn.Name) {
continue
}
fn.FuncNamePrefix = g.FuncNamePrefix
buff.Write([]byte("\n\n"))
if err := funcTpl.Execute(buff, &fn); err != nil {
return err
}
}
return cleanimports.Clean(w, buff.Bytes())
}
// Write writes the generated code in the io.Writer
func (g Generator) Write(wr io.Writer) error {
var buf bytes.Buffer
if err := g.mockTmpl.Execute(&buf, g); err != nil {
return err
}
err := cleanimports.Clean(wr, buf.Bytes())
if err != nil {
err = GenerationError{
Err: err,
Code: buf.Bytes(),
}
}
return err
}
