feat: ai-token-ratelimit support setting global rate limit thresholds for routes​ (#2667)

plugins/wasm-go/extensions/ai-token-ratelimit/config/config_test.go

@@ -0,0 +1,218 @@
+package config
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tidwall/gjson"
+)
+
+func TestParseAiTokenRateLimitConfig(t *testing.T) {
+	tests := []struct {
+		name        string
+		json        string
+		expected    AiTokenRateLimitConfig
+		expectedErr error
+	}{
+		{
+			name:        "MissingRuleName",
+			json:        `{}`,
+			expectedErr: errors.New("missing rule_name in config"),
+		},
+		{
+			name: "GlobalThreshold_InvalidThreshold",
+			json: `{
+				"rule_name": "invalid-threshold",
+				"global_threshold": {
+					"token_per_minute": -100
+				}
+			}`,
+			expectedErr: errors.New("failed to parse global_threshold: 'token_per_minute' must be a positive integer, got -100"),
+		},
+		{
+			name: "GlobalThreshold_QueryPerSecond",
+			json: `{
+				"rule_name": "global-route-limit",
+				"global_threshold": {
+					"token_per_second": 100
+				}
+			}`,
+			expected: AiTokenRateLimitConfig{
+				RuleName: "global-route-limit",
+				GlobalThreshold: &GlobalThreshold{
+					Count:      100,
+					TimeWindow: Second,
+				},
+				RejectedCode: DefaultRejectedCode,
+				RejectedMsg:  DefaultRejectedMsg,
+			},
+		},
+		{
+			name: "GlobalThreshold_QueryPerMinute",
+			json: `{
+				"rule_name": "global-route-limit",
+				"global_threshold": {
+					"token_per_minute": 1000
+				}
+			}`,
+			expected: AiTokenRateLimitConfig{
+				RuleName: "global-route-limit",
+				GlobalThreshold: &GlobalThreshold{
+					Count:      1000,
+					TimeWindow: SecondsPerMinute,
+				},
+				RejectedCode: DefaultRejectedCode,
+				RejectedMsg:  DefaultRejectedMsg,
+			},
+		},
+		{
+			name: "RuleItems_InvalidThreshold",
+			json: `{
+				"rule_name": "invalid-threshold",
+				"rule_items": [
+					{
+						"limit_by_header": "x-test",
+						"limit_keys": [
+							{"key": "key1", "token_per_minute": -100}
+						]
+					}
+				]
+			}`,
+			expectedErr: errors.New("failed to parse rule_item in rule_items: 'token_per_minute' must be a positive integer for key 'key1', got -100"),
+		},
+		{
+			name: "RuleItems_SingleRule",
+			json: `{
+				"rule_name": "rule-based-limit",
+				"rule_items": [
+					{
+						"limit_by_header": "x-test",
+						"limit_keys": [
+							{"key": "key1", "token_per_second": 10}
+						]
+					}
+				]
+			}`,
+			expected: AiTokenRateLimitConfig{
+				RuleName: "rule-based-limit",
+				RuleItems: []LimitRuleItem{
+					{
+						LimitType: LimitByHeaderType,
+						Key:       "x-test",
+						ConfigItems: []LimitConfigItem{
+							{
+								ConfigType: ExactType,
+								Key:        "key1",
+								Count:      10,
+								TimeWindow: Second,
+							},
+						},
+					},
+				},
+				RejectedCode: DefaultRejectedCode,
+				RejectedMsg:  DefaultRejectedMsg,
+			},
+		},
+		{
+			name: "RuleItems_MultipleRules",
+			json: `{
+				"rule_name": "multi-rule-limit",
+				"rule_items": [
+					{
+						"limit_by_param": "user_id",
+						"limit_keys": [
+							{"key": "123", "token_per_hour": 50}
+						]
+					},
+					{
+						"limit_by_per_cookie": "session_id",
+						"limit_keys": [
+							{"key": "*", "token_per_day": 100}
+						]
+					}
+				]
+			}`,
+			expected: AiTokenRateLimitConfig{
+				RuleName: "multi-rule-limit",
+				RuleItems: []LimitRuleItem{
+					{
+						LimitType: LimitByParamType,
+						Key:       "user_id",
+						ConfigItems: []LimitConfigItem{
+							{
+								ConfigType: ExactType,
+								Key:        "123",
+								Count:      50,
+								TimeWindow: SecondsPerHour,
+							},
+						},
+					},
+					{
+						LimitType: LimitByPerCookieType,
+						Key:       "session_id",
+						ConfigItems: []LimitConfigItem{
+							{
+								ConfigType: AllType,
+								Key:        "*",
+								Count:      100,
+								TimeWindow: SecondsPerDay,
+							},
+						},
+					},
+				},
+				RejectedCode: DefaultRejectedCode,
+				RejectedMsg:  DefaultRejectedMsg,
+			},
+		},
+		{
+			name: "Conflict_GlobalThresholdAndRuleItems",
+			json: `{
+				"rule_name": "test-conflict",
+				"global_threshold": {"token_per_second": 100},
+				"rule_items": [{"limit_by_header": "x-test"}]
+			}`,
+			expectedErr: errors.New("'global_threshold' and 'rule_items' cannot be set at the same time"),
+		},
+		{
+			name: "Missing_GlobalThresholdAndRuleItems",
+			json: `{
+				"rule_name": "test-missing"
+			}`,
+			expectedErr: errors.New("at least one of 'global_threshold' or 'rule_items' must be set"),
+		},
+		{
+			name: "Custom_RejectedCodeAndMessage",
+			json: `{
+				"rule_name": "custom-reject",
+				"rejected_code": 403,
+				"rejected_msg": "Forbidden",
+				"global_threshold": {"token_per_second": 100}
+			}`,
+			expected: AiTokenRateLimitConfig{
+				RuleName: "custom-reject",
+				GlobalThreshold: &GlobalThreshold{
+					Count:      100,
+					TimeWindow: Second,
+				},
+				RejectedCode: 403,
+				RejectedMsg:  "Forbidden",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			var config AiTokenRateLimitConfig
+			result := gjson.Parse(tt.json)
+			err := ParseAiTokenRateLimitConfig(result, &config)
+
+			if tt.expectedErr != nil {
+				assert.EqualError(t, err, tt.expectedErr.Error())
+			} else {
+				assert.NoError(t, err)
+				assert.Equal(t, tt.expected, config)
+			}
+		})
+	}
+}