Moving GeneratedContent to the intepreter and renaming to GenerationO…

…utput

* generator.Generate now returns generator.GeneratedEntities

generator/generated_entities.go

@@ -0,0 +1,16 @@
+package generator
+
+type GeneratedEntities []GeneratedFields
+
+type GeneratedFields map[string]interface{}
+
+func NewGeneratedEntities(count int64) GeneratedEntities {
+	return make([]GeneratedFields, count)
+}
+
+func (ge GeneratedEntities) Concat(newEntities GeneratedEntities) GeneratedEntities {
+	for _, entity := range newEntities {
+		ge = append(ge, entity)
+	}
+	return ge
+}

generator/generator.go

@@ -121,8 +121,7 @@ func (g *Generator) WithField(fieldName, fieldType string, fieldOpts interface{}
 	return nil
 }
 
-func (g *Generator) Generate(count int64) GeneratedContent {
-	result := NewGeneratedContent()
+func (g *Generator) Generate(count int64) GeneratedEntities {
 	entities := NewGeneratedEntities(count)
 	for i := int64(0); i < count; i++ {
 		entity := GeneratedFields{}
@@ -131,7 +130,5 @@ func (g *Generator) Generate(count int64) GeneratedContent {
 		}
 		entities[i] = entity
 	}
-	result[g.Name] = entities
-
-	return result
+	return entities
 }

generator/generator_test.go

@@ -32,19 +32,19 @@ func TestExtendGenerator(t *testing.T) {
 	m.WithStaticField("species", "h00man")
 	m.WithStaticField("name", "kyle")
 
-	data := NewGeneratedContent()
+	data := GeneratedEntities{}
 
 	data = g.Generate(1)
 
-	base := data["thing"][0]
+	base := data[0]
 
 	AssertEqual(t, "human", base["species"])
 	AssertEqual(t, 10, len(base["name"].(string)))
 	Assert(t, isBetween(base["age"].(float64), 2, 4), "base entity failed to generate the correct age")
 
 	data = m.Generate(1)
 
-	extended := data["thang"][0]
+	extended := data[0]
 	AssertEqual(t, "h00man", extended["species"])
 	AssertEqual(t, "kyle", extended["name"].(string))
 	Assert(t, isBetween(extended["age"].(float64), 2, 4), "extended entity failed to generate the correct age")
@@ -178,7 +178,7 @@ func TestFieldOptsMatchesFieldType(t *testing.T) {
 }
 
 func TestGenerateProducesGeneratedContent(t *testing.T) {
-	data := NewGeneratedContent()
+	data := GeneratedEntities{}
 	logger := GetLogger(t)
 	g := NewGenerator("thing", logger)
 	timeMin, _ := time.Parse("2006-01-02", "1945-01-01")
@@ -203,7 +203,7 @@ func TestGenerateProducesGeneratedContent(t *testing.T) {
 
 	data = g.Generate(3)
 
-	AssertEqual(t, 3, len(data["thing"]))
+	AssertEqual(t, 3, len(data))
 
 	var testFields = []struct {
 		fieldName string
@@ -228,7 +228,7 @@ func TestGenerateProducesGeneratedContent(t *testing.T) {
 		{"q", uuid.NewV4()},
 	}
 
-	entity := data["thing"][0]
+	entity := data[0]
 	for _, field := range testFields {
 		actual := reflect.TypeOf(entity[field.fieldName])
 		expected := reflect.TypeOf(field.fieldType)

interpreter/interpreter.go

@@ -34,22 +34,22 @@ func debug(format string, tokens ...interface{}) {
 }
 
 type Interpreter struct {
-	entities         map[string]*generator.Generator // TODO: should probably be a more generic symbol table or possibly the parent scope
-	generatedContent generator.GeneratedContent
+	entities map[string]*generator.Generator // TODO: should probably be a more generic symbol table or possibly the parent scope
+	output   GenerationOutput
 }
 
 func New() *Interpreter {
 	return &Interpreter{
-		entities:         make(map[string]*generator.Generator),
-		generatedContent: generator.NewGeneratedContent(),
+		entities: make(map[string]*generator.Generator),
+		output:   GenerationOutput{},
 	}
 }
 
 func (i *Interpreter) WriteGeneratedContent(dest string, filePerEntity bool) error {
 	if filePerEntity {
-		return i.generatedContent.WriteFilePerKey()
+		return i.output.writeFilePerKey()
 	} else {
-		return i.generatedContent.WriteContentToFile(dest)
+		return i.output.writeToFile(dest)
 	}
 }
 
@@ -291,6 +291,6 @@ func (i *Interpreter) GenerateFromNode(generationNode dsl.Node) error {
 		return generationNode.Err("Must generate at least 1 `%s` entity", entityGenerator.Name)
 	}
 
-	i.generatedContent.Append(entityGenerator.Generate(count))
+	i.output.addAndAppend(entityGenerator.Name, entityGenerator.Generate(count))
 	return nil
 }

interpreter/output.go

@@ -0,0 +1,77 @@
+package interpreter
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	g "github.com/ThoughtWorksStudios/datagen/generator"
+	"io"
+	"os"
+)
+
+type GenerationOutput map[string]g.GeneratedEntities
+
+func (output GenerationOutput) addAndAppend(entityName string, entities g.GeneratedEntities) {
+	if _, ok := output[entityName]; ok {
+		output[entityName] = output[entityName].Concat(entities)
+	} else {
+		output[entityName] = entities
+	}
+}
+
+func (output GenerationOutput) writeFilePerKey() error {
+	for k, v := range output {
+		out, err := createWriterFor(fmt.Sprintf("%s.json", k))
+		if err != nil {
+			return err
+		}
+		d := GenerationOutput{}
+		d[k] = v
+		if err = d.write(out); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (output GenerationOutput) writeToFile(dest string) error {
+	out, err := createWriterFor(dest)
+	if err != nil {
+		return err
+	}
+
+	return output.write(out)
+}
+
+func (output GenerationOutput) write(out io.Writer) error {
+	if closeable, doClose := isClosable(out); doClose {
+		defer closeable.Close()
+	}
+
+	writer := bufio.NewWriter(out)
+	encoder := json.NewEncoder(writer)
+	encoder.SetIndent("", "\t")
+
+	if err := encoder.Encode(output); err != nil {
+		return err
+	}
+
+	return writer.Flush()
+}
+
+func isClosable(v interface{}) (io.Closer, bool) {
+	closeable, doClose := v.(io.Closer)
+	return closeable, doClose
+}
+
+func createWriterFor(filename string) (io.Writer, error) {
+	var f interface{}
+	var err error
+	f, err = os.Create(filename)
+	if err != nil {
+		return nil, err
+	}
+	out, _ := f.(io.Writer)
+
+	return out, nil
+}

interpreter/output_test.go

@@ -0,0 +1,30 @@
+package interpreter
+
+import (
+	g "github.com/ThoughtWorksStudios/datagen/generator"
+	. "github.com/ThoughtWorksStudios/datagen/test_helpers"
+	"reflect"
+	"testing"
+)
+
+func TestAppendingToGenerationOutput(t *testing.T) {
+	actual := GenerationOutput{}
+
+	beast := g.GeneratedEntities{g.GeneratedFields{"of the beast": 666}}
+
+	expected := GenerationOutput{"sign": beast}
+	actual.addAndAppend("sign", beast)
+
+	Assert(t, reflect.DeepEqual(expected, actual), "expected \n%v\n to be equal to \n%v\n but wasn't", expected, actual)
+
+	rick := g.GeneratedFields{
+		"of Rick": "wubba lubba dub dub!!!!",
+	}
+
+	expected = GenerationOutput{
+		"sign": append(beast, rick),
+	}
+
+	actual.addAndAppend("sign", g.GeneratedEntities{rick})
+	Assert(t, reflect.DeepEqual(expected, actual), "expected \n%v\n to be equal to \n%v\n but wasn't", expected, actual)
+}