Merge pull request #141 from winebarrel/add_tx_test

Add tx test
winebarrel · Nov 2, 2024 · 5d870d3 · 5d870d3
2 parents e8b844f + f55a60a
commit 5d870d3
Showing 1 changed file with 146 additions and 12 deletions.
diff --git a/task_test.go b/task_test.go
@@ -326,11 +326,11 @@ func TestAcc_Task_MySQLSession(t *testing.T) {
 
 	db, err := sql.Open("mysql", testDSN_MySQL)
 	require.NoError(err)
-	_, err = db.Exec("create table qube_sess_test (data text not null)")
+	_, err = db.Exec("create table qube_test (data text not null)")
 	require.NoError(err)
 
 	t.Cleanup(func() {
-		_, err := db.Exec("drop table qube_sess_test")
+		_, err := db.Exec("drop table qube_test")
 		require.NoError(err)
 		db.Close()
 	})
@@ -339,7 +339,7 @@ func TestAcc_Task_MySQLSession(t *testing.T) {
 	defer os.Remove(f.Name())
 	f.WriteString(`
 {"q":"set @sess_data = 'ABC'"}
-{"q":"insert into qube_sess_test values (@sess_data)"}
+{"q":"insert into qube_test values (@sess_data)"}
 `)
 	f.Sync()
 
@@ -376,10 +376,10 @@ func TestAcc_Task_MySQLSession(t *testing.T) {
 	assert.Equal(0, report.ErrorQueryCount)
 	assert.NotEqual(0, report.AvgQPS)
 
-	var sessData string
-	err = db.QueryRow("select data from qube_sess_test limit 1").Scan(&sessData)
+	var data string
+	err = db.QueryRow("select data from qube_test limit 1").Scan(&data)
 	require.NoError(err)
-	assert.Equal("ABC", sessData)
+	assert.Equal("ABC", data)
 }
 
 func TestAcc_Task_PostgreSQLSession(t *testing.T) {
@@ -392,11 +392,11 @@ func TestAcc_Task_PostgreSQLSession(t *testing.T) {
 
 	db, err := sql.Open("pgx", testDSN_PostgreSQL)
 	require.NoError(err)
-	_, err = db.Exec("create table qube_sess_test (data text not null)")
+	_, err = db.Exec("create table qube_test (data text not null)")
 	require.NoError(err)
 
 	t.Cleanup(func() {
-		_, err := db.Exec("drop table qube_sess_test")
+		_, err := db.Exec("drop table qube_test")
 		require.NoError(err)
 		db.Close()
 	})
@@ -405,7 +405,7 @@ func TestAcc_Task_PostgreSQLSession(t *testing.T) {
 	defer os.Remove(f.Name())
 	f.WriteString(`
 {"q":"select set_config('my.sess_data', 'ABC', false)"}
-{"q":"insert into qube_sess_test values (current_setting('my.sess_data'))"}
+{"q":"insert into qube_test values (current_setting('my.sess_data'))"}
 `)
 	f.Sync()
 
@@ -442,8 +442,142 @@ func TestAcc_Task_PostgreSQLSession(t *testing.T) {
 	assert.Equal(0, report.ErrorQueryCount)
 	assert.NotEqual(0, report.AvgQPS)
 
-	var sessData string
-	err = db.QueryRow("select data from qube_sess_test limit 1").Scan(&sessData)
+	var data string
+	err = db.QueryRow("select data from qube_test limit 1").Scan(&data)
 	require.NoError(err)
-	assert.Equal("ABC", sessData)
+	assert.Equal("ABC", data)
+}
+
+func TestAcc_Task_MySQLTx(t *testing.T) {
+	if !testAcc {
+		t.Skip()
+	}
+
+	assert := assert.New(t)
+	require := require.New(t)
+
+	db, err := sql.Open("mysql", testDSN_MySQL)
+	require.NoError(err)
+	_, err = db.Exec("create table qube_test (data text not null)")
+	require.NoError(err)
+
+	t.Cleanup(func() {
+		_, err := db.Exec("drop table qube_test")
+		require.NoError(err)
+		db.Close()
+	})
+
+	f, _ := os.CreateTemp("", "")
+	defer os.Remove(f.Name())
+	f.WriteString(`
+{"q":"begin"}
+{"q":"insert into qube_test values ('ABC')"}
+{"q":"rollback"}
+`)
+	f.Sync()
+
+	task := &qube.Task{
+		Options: &qube.Options{
+			AgentOptions: qube.AgentOptions{
+				Force: false,
+			},
+			DataOptions: qube.DataOptions{
+				DataFiles:  []string{f.Name()},
+				Key:        "q",
+				Loop:       true,
+				Random:     false,
+				CommitRate: 0,
+			},
+			DBConfig: qube.DBConfig{
+				DSN:    testDSN_MySQL,
+				Driver: qube.DBDriverMySQL,
+				Noop:   false,
+			},
+			Nagents:  10,
+			Rate:     0,
+			Time:     3 * time.Second,
+			Progress: false,
+		},
+		ID: testUUID,
+	}
+
+	report, err := task.Run()
+
+	require.NoError(err)
+	assert.Equal(testUUID, report.ID)
+	assert.NotEqual(0, report.QueryCount)
+	assert.Equal(0, report.ErrorQueryCount)
+	assert.NotEqual(0, report.AvgQPS)
+
+	var cnt int
+	err = db.QueryRow("select count(*) from qube_test").Scan(&cnt)
+	require.NoError(err)
+	assert.Equal(0, cnt)
+}
+
+func TestAcc_Task_PostgreSQLTx(t *testing.T) {
+	if !testAcc {
+		t.Skip()
+	}
+
+	assert := assert.New(t)
+	require := require.New(t)
+
+	db, err := sql.Open("pgx", testDSN_PostgreSQL)
+	require.NoError(err)
+	_, err = db.Exec("create table qube_test (data text not null)")
+	require.NoError(err)
+
+	t.Cleanup(func() {
+		_, err := db.Exec("drop table qube_test")
+		require.NoError(err)
+		db.Close()
+	})
+
+	f, _ := os.CreateTemp("", "")
+	defer os.Remove(f.Name())
+	f.WriteString(`
+{"q":"begin"}
+{"q":"insert into qube_test values ('ABC')"}
+{"q":"rollback"}
+`)
+	f.Sync()
+
+	task := &qube.Task{
+		Options: &qube.Options{
+			AgentOptions: qube.AgentOptions{
+				Force: false,
+			},
+			DataOptions: qube.DataOptions{
+				DataFiles:  []string{f.Name()},
+				Key:        "q",
+				Loop:       true,
+				Random:     false,
+				CommitRate: 0,
+			},
+			DBConfig: qube.DBConfig{
+				DSN:    testDSN_PostgreSQL,
+				Driver: qube.DBDriverPostgreSQL,
+				Noop:   false,
+			},
+			Nagents:  10,
+			Rate:     0,
+			Time:     3 * time.Second,
+			Progress: false,
+		},
+		ID: testUUID,
+	}
+
+	report, err := task.Run()
+
+	require.NoError(err)
+	assert.Equal(testUUID, report.ID)
+	assert.NotEqual(0, report.QueryCount)
+	assert.Equal(0, report.ErrorQueryCount)
+	assert.NotEqual(0, report.AvgQPS)
+
+	var cnt int
+	err = db.QueryRow("select count(*) from qube_test").Scan(&cnt)
+	require.NoError(err)
+	assert.Equal(0, cnt)
 }