feat: support to set the request timeout (#15)

LinuxSuRen · LinuxSuRen · commit d862afa8e554 · 2023-03-30T18:50:33.000+08:00
Co-authored-by: rick &lt;LinuxSuRen@users.noreply.github.com&gt;
diff --git a/cmd/run.go b/cmd/run.go
@@ -9,6 +9,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/linuxsuren/api-testing/pkg/limit"
 	"github.com/linuxsuren/api-testing/pkg/render"
 	"github.com/linuxsuren/api-testing/pkg/runner"
 	"github.com/linuxsuren/api-testing/pkg/testing"
@@ -17,10 +18,16 @@ import (
 )
 
 type runOption struct {
-	pattern  string
-	duration time.Duration
-	thread   int64
-	context  context.Context
+	pattern            string
+	duration           time.Duration
+	requestTimeout     time.Duration
+	requestIgnoreError bool
+	thread             int64
+	context            context.Context
+	qps                int32
+	burst              int32
+	limiter            limit.RateLimiter
+	startTime          time.Time
 }
 
 // CreateRunCommand returns the run command
@@ -40,13 +47,23 @@ See also https://github.com/LinuxSuRen/api-testing/tree/master/sample`,
 	flags.StringVarP(&opt.pattern, "pattern", "p", "test-suite-*.yaml",
 		"The file pattern which try to execute the test cases")
 	flags.DurationVarP(&opt.duration, "duration", "", 0, "Running duration")
+	flags.DurationVarP(&opt.requestTimeout, "request-timeout", "", time.Minute, "Timeout for per request")
+	flags.BoolVarP(&opt.requestIgnoreError, "request-ignore-error", "", false, "Indicate if ignore the request error")
 	flags.Int64VarP(&opt.thread, "thread", "", 1, "Threads of the execution")
+	flags.Int32VarP(&opt.qps, "qps", "", 5, "QPS")
+	flags.Int32VarP(&opt.burst, "burst", "", 5, "burst")
 	return
 }
 
 func (o *runOption) runE(cmd *cobra.Command, args []string) (err error) {
 	var files []string
+	o.startTime = time.Now()
 	o.context = cmd.Context()
+	o.limiter = limit.NewDefaultRateLimiter(o.qps, o.burst)
+	defer func() {
+		cmd.Printf("consume: %s\n", time.Now().Sub(o.startTime).String())
+		o.limiter.Stop()
+	}()
 
 	if files, err = filepath.Glob(o.pattern); err == nil {
 		for i := range files {
@@ -69,13 +86,15 @@ func (o *runOption) runSuiteWithDuration(suite string) (err error) {
 		// make sure having a valid timer
 		timeout = time.NewTicker(time.Second)
 	}
-	errChannel := make(chan error, 10)
+	errChannel := make(chan error, 10*o.thread)
+	stopSingal := make(chan struct{}, 1)
 	var wait sync.WaitGroup
 
 	for !stop {
 		select {
 		case <-timeout.C:
 			stop = true
+			stopSingal <- struct{}{}
 		case err = <-errChannel:
 			if err != nil {
 				stop = true
@@ -85,32 +104,38 @@ func (o *runOption) runSuiteWithDuration(suite string) (err error) {
 				continue
 			}
 			wait.Add(1)
-			if o.duration <= 0 {
-				stop = true
-			}
 
-			go func(ch chan error) {
+			go func(ch chan error, sem *semaphore.Weighted) {
+				now := time.Now()
 				defer sem.Release(1)
 				defer wait.Done()
+				defer func() {
+					fmt.Println("routing end with", time.Now().Sub(now))
+				}()
 
-				ctx := getDefaultContext()
-				ch <- runSuite(suite, ctx)
-			}(errChannel)
+				dataContext := getDefaultContext()
+				ch <- o.runSuite(suite, dataContext, o.context, stopSingal)
+			}(errChannel, sem)
 		}
 	}
-	err = <-errChannel
+
+	select {
+	case err = <-errChannel:
+	case <-stopSingal:
+	}
+
 	wait.Wait()
 	return
 }
 
-func runSuite(suite string, ctx map[string]interface{}) (err error) {
+func (o *runOption) runSuite(suite string, dataContext map[string]interface{}, ctx context.Context, stopSingal chan struct{}) (err error) {
 	var testSuite *testing.TestSuite
 	if testSuite, err = testing.Parse(suite); err != nil {
 		return
 	}
 
 	var result string
-	if result, err = render.Render("base api", testSuite.API, ctx); err == nil {
+	if result, err = render.Render("base api", testSuite.API, dataContext); err == nil {
 		testSuite.API = result
 		testSuite.API = strings.TrimSuffix(testSuite.API, "/")
 	} else {
@@ -123,12 +148,25 @@ func runSuite(suite string, ctx map[string]interface{}) (err error) {
 			testCase.Request.API = fmt.Sprintf("%s%s", testSuite.API, testCase.Request.API)
 		}
 
-		setRelativeDir(suite, &testCase)
 		var output interface{}
-		if output, err = runner.RunTestCase(&testCase, ctx); err != nil {
+		select {
+		case <-stopSingal:
 			return
+		default:
+			// reuse the API prefix
+			if strings.HasPrefix(testCase.Request.API, "/") {
+				testCase.Request.API = fmt.Sprintf("%s%s", testSuite.API, testCase.Request.API)
+			}
+
+			setRelativeDir(suite, &testCase)
+			o.limiter.Accept()
+
+			ctxWithTimeout, _ := context.WithTimeout(ctx, o.requestTimeout)
+			if output, err = runner.RunTestCase(&testCase, dataContext, ctxWithTimeout); err != nil && !o.requestIgnoreError {
+				return
+			}
 		}
-		ctx[testCase.Name] = output
+		dataContext[testCase.Name] = output
 	}
 	return
 }
diff --git a/cmd/run_test.go b/cmd/run_test.go
@@ -1,10 +1,13 @@
 package cmd
 
 import (
+	"context"
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/h2non/gock"
+	"github.com/linuxsuren/api-testing/pkg/limit"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 )
@@ -46,7 +49,13 @@ func TestRunSuite(t *testing.T) {
 
 			tt.prepare()
 			ctx := getDefaultContext()
-			err := runSuite(tt.suiteFile, ctx)
+			opt := &runOption{
+				requestTimeout: 30 * time.Second,
+				limiter:        limit.NewDefaultRateLimiter(0, 0),
+			}
+			stopSingal := make(chan struct{}, 1)
+
+			err := opt.runSuite(tt.suiteFile, ctx, context.TODO(), stopSingal)
 			assert.Equal(t, tt.hasError, err != nil, err)
 		})
 	}
diff --git a/pkg/limit/limiter.go b/pkg/limit/limiter.go
@@ -0,0 +1,95 @@
+package limit
+
+import (
+	"sync"
+	"time"
+)
+
+type RateLimiter interface {
+	TryAccept() bool
+	Accept()
+	Stop()
+	Burst() int32
+}
+
+type defaultRateLimiter struct {
+	qps       int32
+	burst     int32
+	lastToken time.Time
+	singal    chan struct{}
+	mu        sync.Mutex
+}
+
+func NewDefaultRateLimiter(qps, burst int32) RateLimiter {
+	if qps <= 0 {
+		qps = 5
+	}
+	if burst <= 0 {
+		burst = 5
+	}
+	limiter := &defaultRateLimiter{
+		qps:    qps,
+		burst:  burst,
+		singal: make(chan struct{}, 1),
+	}
+	go limiter.updateBurst()
+	return limiter
+}
+
+func (r *defaultRateLimiter) TryAccept() bool {
+	_, ok := r.resver()
+	return ok
+}
+
+func (r *defaultRateLimiter) resver() (delay time.Duration, ok bool) {
+	delay = time.Now().Sub(r.lastToken) / time.Millisecond
+	r.lastToken = time.Now()
+	if delay > 0 {
+		ok = true
+	} else if r.Burst() > 0 {
+		r.Setburst(r.Burst() - 1)
+		ok = true
+	} else {
+		delay = time.Second / time.Duration(r.qps)
+	}
+	return
+}
+
+func (r *defaultRateLimiter) Accept() {
+	delay, ok := r.resver()
+	if ok {
+		return
+	}
+
+	if delay > 0 {
+		time.Sleep(delay)
+	}
+	return
+}
+
+func (r *defaultRateLimiter) Setburst(burst int32) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.burst = burst
+}
+
+func (r *defaultRateLimiter) Burst() int32 {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.burst
+}
+
+func (r *defaultRateLimiter) Stop() {
+	r.singal <- struct{}{}
+}
+
+func (r *defaultRateLimiter) updateBurst() {
+	for {
+		select {
+		case <-time.After(time.Second):
+			r.Setburst(r.Burst() + r.qps)
+		case <-r.singal:
+			return
+		}
+	}
+}
diff --git a/pkg/limit/limiter_test.go b/pkg/limit/limiter_test.go
@@ -0,0 +1,27 @@
+package limit
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestXxx(t *testing.T) {
+	limiter := NewDefaultRateLimiter(1, 1)
+	num := 0
+
+	loop := true
+	go func(l RateLimiter) {
+		for loop {
+			l.Accept()
+			num += 1
+		}
+	}(limiter)
+
+	select {
+	case <-time.After(time.Second):
+		loop = false
+	}
+	assert.True(t, num <= 10)
+}
diff --git a/pkg/runner/simple.go b/pkg/runner/simple.go
@@ -2,6 +2,7 @@ package runner
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -11,7 +12,6 @@ import (
 	"os"
 	"reflect"
 	"strings"
-	"time"
 
 	"github.com/andreyvit/diff"
 	"github.com/antonmedv/expr"
@@ -22,7 +22,7 @@ import (
 )
 
 // RunTestCase runs the test case
-func RunTestCase(testcase *testing.TestCase, ctx interface{}) (output interface{}, err error) {
+func RunTestCase(testcase *testing.TestCase, dataContext interface{}, ctx context.Context) (output interface{}, err error) {
 	fmt.Printf("start to run: '%s'\n", testcase.Name)
 	if err = doPrepare(testcase); err != nil {
 		err = fmt.Errorf("failed to prepare, error: %v", err)
@@ -37,9 +37,7 @@ func RunTestCase(testcase *testing.TestCase, ctx interface{}) (output interface{
 		}
 	}()
 
-	client := http.Client{
-		Timeout: time.Second * 30,
-	}
+	client := http.Client{}
 	var requestBody io.Reader
 	if testcase.Request.Body != "" {
 		requestBody = bytes.NewBufferString(testcase.Request.Body)
@@ -51,7 +49,7 @@ func RunTestCase(testcase *testing.TestCase, ctx interface{}) (output interface{
 		requestBody = bytes.NewBufferString(string(data))
 	}
 
-	if err = testcase.Request.Render(ctx); err != nil {
+	if err = testcase.Request.Render(dataContext); err != nil {
 		return
 	}
 
@@ -76,7 +74,7 @@ func RunTestCase(testcase *testing.TestCase, ctx interface{}) (output interface{
 	}
 
 	var request *http.Request
-	if request, err = http.NewRequest(testcase.Request.Method, testcase.Request.API, requestBody); err != nil {
+	if request, err = http.NewRequestWithContext(ctx, testcase.Request.Method, testcase.Request.API, requestBody); err != nil {
 		return
 	}
 
diff --git a/pkg/runner/simple_test.go b/pkg/runner/simple_test.go
@@ -1,6 +1,7 @@
 package runner
 
 import (
+	"context"
 	"errors"
 	"net/http"
 	"testing"
@@ -365,7 +366,7 @@ func TestTestCase(t *testing.T) {
 			if tt.prepare != nil {
 				tt.prepare()
 			}
-			output, err := RunTestCase(tt.testCase, tt.ctx)
+			output, err := RunTestCase(tt.testCase, tt.ctx, context.TODO())
 			tt.verify(t, output, err)
 		})
 	}
