Merge pull request #133 from winebarrel/support_multi_data

Support multiple input data

.github/workflows/test.yml

@@ -32,7 +32,11 @@ jobs:
       - run: make testacc
       - name: Check race conditions
         run: |
-          echo '{"q":"select 1"}' > data.jsonl
-          go run -race ./cmd/qube -d 'root@tcp(127.0.0.1:13306)/' -f data.jsonl -t 3s -n 10
-          echo '{"q":"invalid"}' >> data.jsonl
-          go run -race ./cmd/qube -d 'postgres://postgres@localhost:15432' -f data.jsonl -t 3s -n 10 --force
+          echo '{"q":"select 1"}' > data1.jsonl
+          echo '{"q":"select 2"}' > data2.jsonl
+          echo '{"q":"select 3"}' > data3.jsonl
+          go run -race ./cmd/qube -d 'root@tcp(127.0.0.1:13306)/' -f data1.jsonl,data2.jsonl,data3.jsonl -t 3s -n 10
+          echo '{"q":"invalid"}' >> data1.jsonl
+          echo '{"q":"invalid"}' >> data2.jsonl
+          echo '{"q":"invalid"}' >> data3.jsonl
+          go run -race ./cmd/qube -d 'postgres://postgres@localhost:15432' -f data1.jsonl,data2.jsonl,data3.jsonl -t 3s -n 10 --force

README.md

@@ -16,12 +16,13 @@ brew install winebarrel/qube/qube
 ## Usage
 
 ```
-Usage: qube --data-file=STRING --dsn=STRING [flags]
+Usage: qube --data-files=DATA-FILES,... --dsn=STRING [flags]
 
 Flags:
   -h, --help                Show help.
       --[no-]force          Do not abort test on error. (default: disabled)
-  -f, --data-file=STRING    NDJSON file path of queries to execute.
+  -f, --data-files=DATA-FILES,...
+                            NDJSON file list of queries to execute.
       --key="q"             Key name of the query field in the test data. e.g.
                             {"q":"SELECT ..."}
       --[no-]loop           Return to the beginning after reading the test data.
@@ -60,7 +61,9 @@ $ qube -d 'root@tcp(127.0.0.1:13306)/' -f data.jsonl -n 5 -t 10s
   "ElapsedTime": "10.001173875s",
   "Options": {
     "Force": false,
-    "DataFile": "data.jsonl",
+    "DataFile": [
+      "data.jsonl"
+    ],
     "Key": "q",
     "Loop": true,
     "Random": false,

agent.go

@@ -34,7 +34,7 @@ func NewAgent(taskID string, agentNum uint64, options *Options, rec *Recorder, l
 		return nil, err
 	}
 
-	data, err := NewData(options)
+	data, err := NewData(options, agentNum)
 
 	if err != nil {
 		return nil, err

agent_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"os"
 	"testing"
 	"time"
@@ -31,7 +32,7 @@ func Test_Agent(t *testing.T) {
 			Force: false,
 		},
 		DataOptions: qube.DataOptions{
-			DataFile:   f.Name(),
+			DataFiles:  []string{f.Name()},
 			Key:        "q",
 			Loop:       true,
 			Random:     false,
@@ -62,5 +63,62 @@ func Test_Agent(t *testing.T) {
 	err = eg.Wait()
 
 	require.True(err == nil || errors.Is(err, context.Canceled))
-	assert.Regexp(`select 1`, buf.String())
+	assert.Regexp("select 1", buf.String())
+}
+
+func Test_Agent_MultiData(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	f1, _ := os.CreateTemp("", "")
+	defer os.Remove(f1.Name())
+	f1.WriteString(`{"q":"select 1"}` + "\n") //nolint:errcheck
+	f1.Sync()                                 //nolint:errcheck
+
+	f2, _ := os.CreateTemp("", "")
+	defer os.Remove(f2.Name())
+	f2.WriteString(`{"q":"select 2"}` + "\n") //nolint:errcheck
+	f2.Sync()                                 //nolint:errcheck
+
+	options := &qube.Options{
+		AgentOptions: qube.AgentOptions{
+			Force: false,
+		},
+		DataOptions: qube.DataOptions{
+			DataFiles:  []string{f1.Name(), f2.Name()},
+			Key:        "q",
+			Loop:       true,
+			Random:     false,
+			CommitRate: 0,
+		},
+		DBConfig: qube.DBConfig{
+			DSN:    testDSN_MySQL,
+			Driver: qube.DBDriverMySQL,
+			Noop:   true,
+		},
+		Nagents:  4,
+		Rate:     0,
+		Time:     1 * time.Second,
+		Progress: false,
+	}
+
+	rec := qube.NewRecorder(testUUID, options)
+	limiter := rate.NewLimiter(rate.Limit(1), 1)
+
+	for i := range 4 {
+		var buf bytes.Buffer
+		options.DBConfig.NullDBOut = &buf
+		agent, err := qube.NewAgent(testUUID, uint64(i), options, rec, limiter)
+		require.NoError(err)
+
+		eg, ctx := errgroup.WithContext(context.Background())
+		ctx, cancel := context.WithCancel(ctx)
+		eg.Go(func() error { return agent.Start(ctx) })
+		time.Sleep(2 * time.Second)
+		cancel()
+		err = eg.Wait()
+
+		require.True(err == nil || errors.Is(err, context.Canceled))
+		assert.Regexp(fmt.Sprintf("select %d", i%2+1), buf.String())
+	}
 }

data.go

@@ -17,11 +17,11 @@ var (
 )
 
 type DataOptions struct {
-	DataFile   string `kong:"short='f',required,help='NDJSON file path of queries to execute.'"`
-	Key        string `kong:"default='q',help='Key name of the query field in the test data. e.g. {\"q\":\"SELECT ...\"}'"`
-	Loop       bool   `kong:"negatable,default='true',help='Return to the beginning after reading the test data. (default: enabled)'"`
-	Random     bool   `kong:"negatable,default='false',help='Randomize the starting position of the test data. (default: disabled)'"`
-	CommitRate uint   `kong:"help='Number of queries to execute \"COMMIT\".'"`
+	DataFiles  []string `kong:"short='f',required,help='NDJSON file list of queries to execute.'"`
+	Key        string   `kong:"default='q',help='Key name of the query field in the test data. e.g. {\"q\":\"SELECT ...\"}'"`
+	Loop       bool     `kong:"negatable,default='true',help='Return to the beginning after reading the test data. (default: enabled)'"`
+	Random     bool     `kong:"negatable,default='false',help='Randomize the starting position of the test data. (default: disabled)'"`
+	CommitRate uint     `kong:"help='Number of queries to execute \"COMMIT\".'"`
 }
 
 type Data struct {
@@ -32,11 +32,12 @@ type Data struct {
 	inTxn  bool
 }
 
-func NewData(options *Options) (*Data, error) {
-	file, err := os.OpenFile(options.DataFile, os.O_RDONLY, 0)
+func NewData(options *Options, agentNum uint64) (*Data, error) {
+	dataFile := options.DataFiles[agentNum%uint64(len(options.DataFiles))]
+	file, err := os.OpenFile(dataFile, os.O_RDONLY, 0)
 
 	if err != nil {
-		return nil, fmt.Errorf("failed to open test data - %s (%w)", options.DataFile, err)
+		return nil, fmt.Errorf("failed to open test data - %s (%w)", dataFile, err)
 	}
 
 	if options.Random {

data_test.go

@@ -1,6 +1,7 @@
 package qube_test
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
@@ -20,12 +21,12 @@ func Test_Data(t *testing.T) {
 
 	options := &qube.Options{
 		DataOptions: qube.DataOptions{
-			DataFile: f.Name(),
-			Key:      "q",
+			DataFiles: []string{f.Name()},
+			Key:       "q",
 		},
 	}
 
-	data, err := qube.NewData(options)
+	data, err := qube.NewData(options, 0)
 	require.NoError(err)
 	defer data.Close()
 
@@ -47,13 +48,13 @@ func Test_Data_Loop(t *testing.T) {
 
 	options := &qube.Options{
 		DataOptions: qube.DataOptions{
-			DataFile: f.Name(),
-			Key:      "q",
-			Loop:     true,
+			DataFiles: []string{f.Name()},
+			Key:       "q",
+			Loop:      true,
 		},
 	}
 
-	data, err := qube.NewData(options)
+	data, err := qube.NewData(options, 0)
 	require.NoError(err)
 	defer data.Close()
 
@@ -76,14 +77,14 @@ func Test_Data_Random(t *testing.T) {
 
 	options := &qube.Options{
 		DataOptions: qube.DataOptions{
-			DataFile: f.Name(),
-			Key:      "q",
-			Loop:     true,
-			Random:   true,
+			DataFiles: []string{f.Name()},
+			Key:       "q",
+			Loop:      true,
+			Random:    true,
 		},
 	}
 
-	data, err := qube.NewData(options)
+	data, err := qube.NewData(options, 0)
 	require.NoError(err)
 	defer data.Close()
 
@@ -106,14 +107,14 @@ func Test_Data_WithCommitRate(t *testing.T) {
 
 	options := &qube.Options{
 		DataOptions: qube.DataOptions{
-			DataFile:   f.Name(),
+			DataFiles:  []string{f.Name()},
 			Key:        "q",
 			Loop:       true,
 			CommitRate: 2,
 		},
 	}
 
-	data, err := qube.NewData(options)
+	data, err := qube.NewData(options, 0)
 	require.NoError(err)
 	defer data.Close()
 
@@ -146,15 +147,54 @@ func Test_Data_WithoutKey(t *testing.T) {
 
 	options := &qube.Options{
 		DataOptions: qube.DataOptions{
-			DataFile: f.Name(),
-			Key:      "q",
-			Loop:     true,
+			DataFiles: []string{f.Name()},
+			Key:       "q",
+			Loop:      true,
 		},
 	}
 
-	data, err := qube.NewData(options)
+	data, err := qube.NewData(options, 0)
 	require.NoError(err)
 
 	_, err = data.Next()
 	assert.ErrorContains(err, `failed to get query field "q" from '{"_q_":"select 1"}'`)
 }
+
+func Test_MultiData(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	f1, _ := os.CreateTemp("", "")
+	defer os.Remove(f1.Name())
+	f1.WriteString(`{"q":"select 1"}` + "\n") //nolint:errcheck
+	f1.Sync()                                 //nolint:errcheck
+
+	f2, _ := os.CreateTemp("", "")
+	defer os.Remove(f2.Name())
+	f2.WriteString(`{"q":"select 2"}` + "\n") //nolint:errcheck
+	f2.Sync()                                 //nolint:errcheck
+
+	f3, _ := os.CreateTemp("", "")
+	defer os.Remove(f3.Name())
+	f3.WriteString(`{"q":"select 3"}` + "\n") //nolint:errcheck
+	f3.Sync()                                 //nolint:errcheck
+
+	options := &qube.Options{
+		DataOptions: qube.DataOptions{
+			DataFiles: []string{f1.Name(), f2.Name(), f3.Name()},
+			Key:       "q",
+		},
+	}
+
+	for i := range 6 {
+		data, err := qube.NewData(options, uint64(i))
+		require.NoError(err)
+		defer data.Close()
+
+		q, err := data.Next()
+		require.NoError(err)
+		assert.Equal(fmt.Sprintf("select %d", i%3+1), q)
+		_, err = data.Next()
+		assert.ErrorIs(err, qube.EOD)
+	}
+}

recorder_test.go

@@ -16,7 +16,7 @@ func Test_Recorder(t *testing.T) {
 			Force: true,
 		},
 		DataOptions: qube.DataOptions{
-			DataFile:   "data.jsonl",
+			DataFiles:  []string{"data.jsonl"},
 			Key:        "q",
 			Loop:       true,
 			Random:     false,
@@ -78,7 +78,9 @@ func Test_Recorder(t *testing.T) {
   "ElapsedTime": "0s",
   "Options": {
     "Force": true,
-    "DataFile": "data.jsonl",
+    "DataFiles": [
+      "data.jsonl"
+    ],
     "Key": "q",
     "Loop": true,
     "Random": false,
@@ -136,7 +138,7 @@ func Test_Recorder_WithError(t *testing.T) {
 			Force: true,
 		},
 		DataOptions: qube.DataOptions{
-			DataFile:   "data.jsonl",
+			DataFiles:  []string{"data.jsonl"},
 			Key:        "q",
 			Loop:       true,
 			Random:     false,
@@ -202,7 +204,9 @@ func Test_Recorder_WithError(t *testing.T) {
   "ElapsedTime": "0s",
   "Options": {
     "Force": true,
-    "DataFile": "data.jsonl",
+    "DataFiles": [
+      "data.jsonl"
+    ],
     "Key": "q",
     "Loop": true,
     "Random": false,