github.com/polarismesh/polaris@v1.17.8/plugin/ratelimit/token/invoke_test.go (about)

     1  /**
     2   * Tencent is pleased to support the open source community by making Polaris available.
     3   *
     4   * Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.
     5   *
     6   * Licensed under the BSD 3-Clause License (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   * https://opensource.org/licenses/BSD-3-Clause
    11   *
    12   * Unless required by applicable law or agreed to in writing, software distributed
    13   * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
    14   * CONDITIONS OF ANY KIND, either express or implied. See the License for the
    15   * specific language governing permissions and limitations under the License.
    16   */
    17  
    18  package token
    19  
    20  import (
    21  	"testing"
    22  
    23  	. "github.com/smartystreets/goconvey/convey"
    24  
    25  	"github.com/polarismesh/polaris/plugin"
    26  )
    27  
    28  // baseConfigOption 返回一个基础的正常option配置
    29  func baseConfigOption() map[string]interface{} {
    30  	return map[string]interface{}{
    31  		"ip-limit": &ResourceLimitConfig{
    32  			Open:                   true,
    33  			Global:                 &BucketRatelimit{true, 10, 2},
    34  			MaxResourceCacheAmount: 100,
    35  		},
    36  		"api-limit": &APILimitConfig{
    37  			Open: true,
    38  			Rules: []*RateLimitRule{{
    39  				Name:  "rule-1",
    40  				Limit: &BucketRatelimit{true, 5, 1},
    41  			}},
    42  			Apis: []*APILimitInfo{{Name: "api-1", Rule: "rule-1"}},
    43  		},
    44  	}
    45  }
    46  
    47  // TestTokenBucket_Name 对插件名字接口实现的测试
    48  func TestTokenBucket_Name(t *testing.T) {
    49  	tb := &tokenBucket{}
    50  	Convey("返回插件名服务预期", t, func() {
    51  		So(tb.Name(), ShouldEqual, PluginName)
    52  	})
    53  }
    54  
    55  // TestTokenBucket_Initialize 测试初始化函数
    56  func TestTokenBucket_Initialize(t *testing.T) {
    57  	configEntry := &plugin.ConfigEntry{Name: PluginName}
    58  	tb := &tokenBucket{}
    59  	Convey("配置option为空，返回失败", t, func() {
    60  		So(tb.Initialize(configEntry), ShouldNotBeNil)
    61  	})
    62  	Convey("配置字段不对，不影响解析，配置不生效", t, func() {
    63  		configEntry.Option = map[string]interface{}{"aaa": 123}
    64  		So(tb.Initialize(configEntry), ShouldBeNil)
    65  	})
    66  	Convey("无效ip-limit配置，返回失败", t, func() {
    67  		configEntry.Option = map[string]interface{}{
    68  			"ip-limit": &ResourceLimitConfig{
    69  				Open:                   true,
    70  				Global:                 nil,
    71  				MaxResourceCacheAmount: 100,
    72  			},
    73  		}
    74  		So(tb.Initialize(configEntry), ShouldNotBeNil)
    75  	})
    76  	Convey("无效api-limit配置，返回失败", t, func() {
    77  		configEntry.Option = map[string]interface{}{
    78  			"api-limit": &APILimitConfig{Open: true},
    79  		}
    80  		So(tb.Initialize(configEntry), ShouldNotBeNil)
    81  	})
    82  	Convey("配置有效，可以正常初始化", t, func() {
    83  		configEntry.Option = baseConfigOption()
    84  		So(tb.Initialize(configEntry), ShouldBeNil)
    85  		So(tb.limiters[plugin.IPRatelimit], ShouldNotBeNil)
    86  		So(tb.limiters[plugin.APIRatelimit], ShouldNotBeNil)
    87  	})
    88  }
    89  
    90  // TestTokenBucket_Allow 测试Allow函数
    91  func TestTokenBucket_Allow(t *testing.T) {
    92  	configEntry := &plugin.ConfigEntry{Name: PluginName}
    93  	configEntry.Option = baseConfigOption()
    94  	tb := &tokenBucket{}
    95  	if err := tb.Initialize(configEntry); err != nil {
    96  		t.Fatalf("error: %s", err.Error())
    97  	}
    98  	ipLimiter := tb.limiters[plugin.IPRatelimit].(*resourceRatelimit)
    99  	apiLimiter := tb.limiters[plugin.APIRatelimit].(*apiRatelimit)
   100  	Convey("IP正常限流", t, func() {
   101  		cnt := 0
   102  		for i := 0; i < ipLimiter.config.Global.Bucket*2; i++ {
   103  			if ok := tb.Allow(plugin.IPRatelimit, "1.2.3.4"); ok {
   104  				cnt++
   105  			}
   106  		}
   107  		So(cnt, ShouldEqual, ipLimiter.config.Global.Bucket)
   108  		// 其他IP可以正常通过
   109  		So(tb.Allow(plugin.IPRatelimit, "2.3.4.5"), ShouldEqual, true)
   110  	})
   111  	Convey("api正常限流", t, func() {
   112  		cnt := 0
   113  		for i := 0; i < apiLimiter.rules["rule-1"].Bucket*2; i++ {
   114  			if ok := tb.Allow(plugin.APIRatelimit, "api-1"); ok {
   115  				cnt++
   116  			}
   117  		}
   118  		So(cnt, ShouldEqual, apiLimiter.rules["rule-1"].Bucket)
   119  		// 其他接口没有限流的可以通过
   120  		So(tb.Allow(plugin.APIRatelimit, "api-2"), ShouldEqual, true)
   121  	})
   122  	Convey("空的key，正常限流", t, func() {
   123  		So(tb.Allow(plugin.APIRatelimit, ""), ShouldEqual, true)
   124  	})
   125  	Convey("非法的限制类型，直接通过", t, func() {
   126  		So(tb.Allow(plugin.RatelimitType(100), "123"), ShouldEqual, true)
   127  	})
   128  }