Unit test for api_collect

go.mod

@@ -3,7 +3,7 @@ module github.com/merico-dev/lake
 go 1.17
 
 require (
-	github.com/agiledragon/gomonkey v2.0.2+incompatible
+	github.com/agiledragon/gomonkey/v2 v2.7.0
 	github.com/gin-contrib/cors v1.3.1
 	github.com/gin-gonic/gin v1.7.4
 	github.com/go-git/go-git/v5 v5.4.2

go.sum

@@ -49,8 +49,8 @@ github.com/ProtonMail/go-crypto v0.0.0-20210428141323-04723f9f07d7 h1:YoJbenK9C6
 github.com/ProtonMail/go-crypto v0.0.0-20210428141323-04723f9f07d7/go.mod h1:z4/9nQmJSSwwds7ejkxaJwO37dru3geImFUdJlaLzQo=
 github.com/acomagu/bufpipe v1.0.3 h1:fxAGrHZTgQ9w5QqVItgzwj235/uYZYgbXitB+dLupOk=
 github.com/acomagu/bufpipe v1.0.3/go.mod h1:mxdxdup/WdsKVreO5GpW4+M/1CE2sMG4jeGJ2sYmHc4=
-github.com/agiledragon/gomonkey v2.0.2+incompatible h1:eXKi9/piiC3cjJD1658mEE2o3NjkJ5vDLgYjCQu0Xlw=
-github.com/agiledragon/gomonkey v2.0.2+incompatible/go.mod h1:2NGfXu1a80LLr2cmWXGBDaHEjb1idR6+FVlX5T3D9hw=
+github.com/agiledragon/gomonkey/v2 v2.7.0 h1:CFT/xdr6xbsIN04Yll4OhKq/vPm0MVD8ykV99jDBesM=
+github.com/agiledragon/gomonkey/v2 v2.7.0/go.mod h1:ap1AmDzcVOAz1YpeJ3TCzIgstoaWLA6jbbgxfB4w2iY=
 github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239 h1:kFOfPq6dUM1hTo4JG6LR5AXSUEsOjtdm0kw0FtQtMJA=
 github.com/anmitsu/go-shlex v0.0.0-20161002113705-648efa622239/go.mod h1:2FmKhYUyUczH0OGQWaF5ceTx0UBShxjsH6f8oGKYe2c=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=

plugins/helper/api_async_client_test.go

@@ -13,7 +13,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/agiledragon/gomonkey"
+	"github.com/agiledragon/gomonkey/v2"
 	"github.com/merico-dev/lake/plugins/core"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"

plugins/helper/api_collector.go

@@ -137,6 +137,9 @@ func (collector *ApiCollector) Execute() error {
 		defer iterator.Close()
 		// throttle input process speed so it can be canceled, create a channel to represent available slots
 		slots := int(math.Ceil(collector.args.ApiClient.GetQps())) * 2
+		if slots <= 0 {
+			return fmt.Errorf("RateLimit can't use the 0 Qps")
+		}
 		slotsChan := make(chan bool, slots)
 		defer close(slotsChan)
 		for i := 0; i < slots; i++ {
@@ -148,17 +151,19 @@ func (collector *ApiCollector) Execute() error {
 
 		var wg sync.WaitGroup
 		ctx := collector.args.Ctx.GetContext()
+
+	out:
 		for iterator.HasNext() {
 			select {
 			// canceled by user, stop
 			case <-ctx.Done():
 				err = ctx.Err()
-				break
+				break out
 			// obtain a slot
 			case <-slotsChan:
 				input, err := iterator.Fetch()
 				if err != nil {
-					break
+					break out
 				}
 				wg.Add(1)
 				go func() {
@@ -176,7 +181,7 @@ func (collector *ApiCollector) Execute() error {
 					}
 				}()
 			case err = <-errors:
-				break
+				break out
 			}
 		}
 		if err == nil {
@@ -425,4 +430,24 @@ func GetRawMessageDirectFromResponse(res *http.Response) ([]json.RawMessage, err
 	return []json.RawMessage{body}, nil
 }
 
+func GetRawMessageArrayFromResponse(res *http.Response) ([]json.RawMessage, error) {
+	rawMessages := []json.RawMessage{}
+
+	if res == nil {
+		return nil, fmt.Errorf("res is nil")
+	}
+	defer res.Body.Close()
+	resBody, err := ioutil.ReadAll(res.Body)
+	if err != nil {
+		return nil, fmt.Errorf("%w %s", err, res.Request.URL.String())
+	}
+
+	err = json.Unmarshal(resBody, &rawMessages)
+	if err != nil {
+		return nil, fmt.Errorf("%w %s %s", err, res.Request.URL.String(), string(resBody))
+	}
+
+	return rawMessages, nil
+}
+
 var _ core.SubTask = (*ApiCollector)(nil)