From 7f309d5eb31c89c45cbaeba5fc119ad8b64893f9 Mon Sep 17 00:00:00 2001
From: justxuewei <justxuewei@apache.org>
Date: Thu, 13 Jan 2022 11:14:59 +0800
Subject: [PATCH] feat(tester): support duration string

---
 pkg/tester/stress.go      | 16 ++++++++++++++--
 pkg/tester/stress_test.go |  2 +-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/pkg/tester/stress.go b/pkg/tester/stress.go
index fcacb58..3bb3800 100644
--- a/pkg/tester/stress.go
+++ b/pkg/tester/stress.go
@@ -38,6 +38,7 @@ func (t *StressTester) Run() {
 	defer func() {
 		now := time.Now()
 		t.endedTime = &now
+		t.log("The stress tester finished, %d requests are sent", t.GetTransactionNum())
 	}()
 
 	if t.startedTime != nil || t.transactionNum != nil || t.endedTime != nil {
@@ -137,6 +138,13 @@ func (t *StressTester) GetAverageRTSeconds() float64 {
 	return t.rt.Load() / float64(t.transactionNum.Load())
 }
 
+func (t *StressTester) GetTransactionNum() int32 {
+	if t.transactionNum == nil {
+		return -1
+	}
+	return t.transactionNum.Load()
+}
+
 func (t *StressTester) SetVerbose(v bool) *StressTester {
 	t.verbose = v
 	return t
@@ -147,8 +155,12 @@ func (t *StressTester) SetUserNum(n int) *StressTester {
 	return t
 }
 
-func (t *StressTester) SetDuration(d time.Duration) *StressTester {
-	t.duration = d
+func (t *StressTester) SetDuration(d string) *StressTester {
+	duration, err := time.ParseDuration(d)
+	if err != nil {
+		panic(nil)
+	}
+	t.duration = duration
 	return t
 }
 
diff --git a/pkg/tester/stress_test.go b/pkg/tester/stress_test.go
index 3ff52eb..cb4ee5a 100644
--- a/pkg/tester/stress_test.go
+++ b/pkg/tester/stress_test.go
@@ -17,7 +17,7 @@ func TestStressTester(t *testing.T) {
 		tester.
 			SetVerbose(true).
 			SetTPS(10).
-			SetDuration(10 * time.Second).
+			SetDuration("10s").
 			SetUserNum(2).
 			SetTestFn(testFn).
 			Run()