file/stdout output: add tests

feliixx · Feb 8, 2022 · 7d48c80 · 7d48c80
1 parent 10a0ddf
commit 7d48c80
Show file tree

Hide file tree

Showing 3 changed files with 177 additions and 2 deletions.
diff --git a/datagen/file_writer.go b/datagen/file_writer.go
@@ -3,6 +3,7 @@ package datagen
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -61,6 +62,11 @@ func (w *fileWriter) write(collections []Collection, seed uint64) (err error) {
 		if err != nil {
 			return fmt.Errorf("fail to create DocumentGenerator for collection '%s'\n%v", collections[i].Name, err)
 		}
+
+		aggs, err := ci.NewAggregatorSlice(collections[i].Content)
+		if len(aggs) > 0 || err != nil {
+			return errors.New("Aggregators are not supported for stdout or file output")
+		}
 	}
 
 	for i := 0; i < len(collections); i++ {
@@ -105,14 +111,14 @@ func (w *fileWriter) writeToStdout(wg *sync.WaitGroup, coll *Collection, tasks <
 
 	buffer := bytes.NewBuffer(make([]byte, 0, 64000))
 
-	// format is : 
+	// format is :
 	//
 	// {
 	//   "dbName.collectionName": [
 	//	   {...},
 	//	   {...},
 	//	   {...}
-	//   ]  	
+	//   ]
 	// }
 	buffer.WriteByte('{')
 	buffer.WriteByte('\n')

diff --git a/datagen/generate_test.go b/datagen/generate_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"log"
 	"os"
@@ -231,6 +232,120 @@ invalid generator for field 'invalid-aggregator'
 	}
 }
 
+func TestJSONOutputToStdout(t *testing.T) {
+	opts := datagen.Options{
+		Configuration: datagen.Configuration{
+			ConfigFile: "testdata/only_id.json",
+			BatchSize:  1000,
+			Output:     "stdout",
+			Seed:       123456789,
+		},
+	}
+
+	var buffer bytes.Buffer
+	err := datagen.Generate(&opts, &buffer)
+	if err != nil {
+		t.Errorf("fail to write to stdout: %v", err)
+	}
+
+	if want, got := "", buffer.String(); want != got {
+		t.Errorf("No logs should be written with output=stdout, but got %s", got)
+	}
+}
+
+func TestJSONOutputToFile(t *testing.T) {
+
+	outputFileName := "/tmp/id.json"
+
+	defer os.Remove(outputFileName)
+
+	opts := datagen.Options{
+		Configuration: datagen.Configuration{
+			ConfigFile: "testdata/only_id.json",
+			BatchSize:  1000,
+			Output:     outputFileName,
+			Seed:       123456789,
+		},
+	}
+
+	var buffer bytes.Buffer
+	err := datagen.Generate(&opts, &buffer)
+	if err != nil {
+		t.Errorf("fail to write to file: %v", err)
+	}
+
+	logs := buffer.String()
+	if !strings.HasPrefix(logs, "Using seed: 123456789") {
+		t.Error("Logs should not be empty")
+	}
+
+	want := `{
+  "mgodatagen_test.test": [
+    {"_id": {"$numberInt":"0"}},
+    {"_id": {"$numberInt":"1"}},
+    {"_id": {"$numberInt":"2"}}
+  ]
+}`
+	got, err := os.ReadFile(outputFileName)
+	if err != nil {
+		t.Errorf("fail to read from %s: %v", outputFileName, err)
+	}
+
+	if want != string(got) {
+		t.Errorf("expected\n\n'%s'\n\nbut got\n\n'%s'", want, got)
+	}
+}
+func TestIndentJSONOutputToFile(t *testing.T) {
+
+	outputFileName := "/tmp/id_indented.json"
+
+	defer os.Remove(outputFileName)
+
+	opts := datagen.Options{
+		Configuration: datagen.Configuration{
+			ConfigFile:  "testdata/only_id.json",
+			BatchSize:   100,
+			Output:      outputFileName,
+			PrettyPrint: true,
+			Seed:        123456789,
+		},
+	}
+
+	err := datagen.Generate(&opts, io.Discard)
+	if err != nil {
+		t.Errorf("fail to write to file: %v", err)
+	}
+
+	want := `{
+  "mgodatagen_test.test": [
+    {
+      "_id": {
+        "$numberInt": "0"
+      }
+    },
+    {
+      "_id": {
+        "$numberInt": "1"
+      }
+    },
+    {
+      "_id": {
+        "$numberInt": "2"
+      }
+    }
+  ]
+}`
+
+	got, err := os.ReadFile(outputFileName)
+	if err != nil {
+		t.Errorf("fail to read from %s: %v", outputFileName, err)
+	}
+
+	if want != string(got) {
+		t.Errorf("expected\n\n'%s'\n\nbut got\n\n'%s'", want, got)
+	}
+}
+
 func TestCollectionContent(t *testing.T) {
 
 	configFile := "generators/testdata/full-bson.json"
@@ -713,6 +828,8 @@ func TestCollectionWithIndexes(t *testing.T) {
 
 func TestGenerate(t *testing.T) {
 
+	defer os.Remove("/tmp/output.json")
+
 	type generateTest struct {
 		name          string
 		options       datagen.Options
@@ -904,6 +1021,45 @@ func TestGenerate(t *testing.T) {
 			correct:     false,
 			errMsgRegex: regexp.MustCompile("^-i | --indexonly and -x | --indexfirst can't be present at the same time. Try to remove the -x | --indexfirst flag.*"),
 		},
+		{
+			name: "stdout output",
+			options: datagen.Options{
+				Configuration: datagen.Configuration{
+					ConfigFile: "testdata/only_id.json",
+					BatchSize:  1000,
+					Output:     "stdout",
+				},
+				General: defaultGeneralOpts,
+			},
+			correct:       true,
+			expectedNbDoc: 0,
+		},
+		{
+			name: "json file output",
+			options: datagen.Options{
+				Configuration: datagen.Configuration{
+					ConfigFile: "testdata/only_id.json",
+					BatchSize:  1000,
+					Output:     "/tmp/output.json",
+				},
+				General: defaultGeneralOpts,
+			},
+			correct:       true,
+			expectedNbDoc: 0,
+		},
+		{
+			name: "stdout output with aggregators",
+			options: datagen.Options{
+				Configuration: datagen.Configuration{
+					ConfigFile: "generators/testdata/full-aggregation.json",
+					BatchSize:  1000,
+					Output:     "stdout",
+				},
+				General: defaultGeneralOpts,
+			},
+			correct:     false,
+			errMsgRegex: regexp.MustCompile("^Aggregators are not supported for stdout or file output.*"),
+		},
 	}
 
 	shardedGenerateTests := []struct {

diff --git a/datagen/testdata/only_id.json b/datagen/testdata/only_id.json
@@ -0,0 +1,13 @@
+[
+  {
+    "database": "mgodatagen_test",
+    "collection": "test",
+    "count": 3,
+    "content": {
+      "_id": {
+        "type": "autoincrement",
+        "autoType": "int"
+      }
+    }
+  }
+]