Refine flow.Rule API

alibaba · Sep 5, 2020 · fbf4a3f · fbf4a3f
1 parent 484a518
commit fbf4a3f
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 5 deletions.
diff --git a/core/flow/rule_manager.go b/core/flow/rule_manager.go
@@ -102,13 +102,46 @@ func LoadRules(rules []*Rule) (bool, error) {
 	return true, err
 }
 
-func GetRules() []*Rule {
+func getRules() []*Rule {
 	tcMux.RLock()
 	defer tcMux.RUnlock()
 
 	return rulesFrom(tcMap)
 }
 
+func getResRules(res string) []*Rule {
+	tcMux.RLock()
+	defer tcMux.RUnlock()
+
+	resTcs, exist := tcMap[res]
+	if !exist {
+		return nil
+	}
+	ret := make([]*Rule, 0, len(resTcs))
+	for _, tc := range resTcs {
+		ret = append(ret, tc.Rule())
+	}
+	return ret
+}
+
+func GetRules() []Rule {
+	rules := getRules()
+	ret := make([]Rule, 0, len(rules))
+	for _, rule := range rules {
+		ret = append(ret, *rule)
+	}
+	return ret
+}
+
+func GetResRules(res string) []Rule {
+	rules := getResRules(res)
+	ret := make([]Rule, 0, len(rules))
+	for _, rule := range rules {
+		ret = append(ret, *rule)
+	}
+	return ret
+}
+
 func ClearRules() error {
 	_, err := LoadRules(nil)
 	return err

diff --git a/core/flow/rule_manager_test.go b/core/flow/rule_manager_test.go
@@ -66,8 +66,110 @@ func TestIsValidFlowRule(t *testing.T) {
 	assert.NoError(t, IsValidRule(goodRule1))
 }
 
-	assert.Error(t, IsValidFlowRule(badRule1))
-	assert.Error(t, IsValidFlowRule(badRule2))
-	assert.Error(t, IsValidFlowRule(badRule3))
-	assert.NoError(t, IsValidFlowRule(goodRule1))
+func TestGetRules(t *testing.T) {
+	t.Run("GetRules", func(t *testing.T) {
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+		r1 := &Rule{
+			ID:               0,
+			Resource:         "abc1",
+			LimitOrigin:      "default",
+			MetricType:       0,
+			Count:            0,
+			RelationStrategy: 0,
+			ControlStrategy: ControlStrategy{
+				AllocationStrategy: Default,
+				ControlEffect:      Reject,
+			},
+			RefResource:       "",
+			WarmUpPeriodSec:   0,
+			MaxQueueingTimeMs: 0,
+		}
+		r2 := &Rule{
+			ID:               1,
+			Resource:         "abc2",
+			LimitOrigin:      "default",
+			MetricType:       0,
+			Count:            0,
+			RelationStrategy: 0,
+			ControlStrategy: ControlStrategy{
+				AllocationStrategy: Default,
+				ControlEffect:      Throttling,
+			},
+			RefResource:       "",
+			WarmUpPeriodSec:   0,
+			MaxQueueingTimeMs: 0,
+		}
+		if _, err := LoadRules([]*Rule{r1, r2}); err != nil {
+			t.Fatal(err)
+		}
+
+		rs1 := GetRules()
+		if rs1[0].Resource == "abc1" {
+			assert.True(t, &rs1[0] != r1)
+			assert.True(t, &rs1[1] != r2)
+			assert.True(t, reflect.DeepEqual(&rs1[0], r1))
+			assert.True(t, reflect.DeepEqual(&rs1[1], r2))
+		} else {
+			assert.True(t, &rs1[0] != r2)
+			assert.True(t, &rs1[1] != r1)
+			assert.True(t, reflect.DeepEqual(&rs1[0], r2))
+			assert.True(t, reflect.DeepEqual(&rs1[1], r1))
+		}
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+	})
+
+	t.Run("getRules", func(t *testing.T) {
+		r1 := &Rule{
+			ID:               0,
+			Resource:         "abc1",
+			LimitOrigin:      "default",
+			MetricType:       0,
+			Count:            0,
+			RelationStrategy: 0,
+			ControlStrategy: ControlStrategy{
+				AllocationStrategy: Default,
+				ControlEffect:      Reject,
+			},
+			RefResource:       "",
+			WarmUpPeriodSec:   0,
+			MaxQueueingTimeMs: 0,
+		}
+		r2 := &Rule{
+			ID:               1,
+			Resource:         "abc2",
+			LimitOrigin:      "default",
+			MetricType:       0,
+			Count:            0,
+			RelationStrategy: 0,
+			ControlStrategy: ControlStrategy{
+				AllocationStrategy: Default,
+				ControlEffect:      Throttling,
+			},
+			RefResource:       "",
+			WarmUpPeriodSec:   0,
+			MaxQueueingTimeMs: 0,
+		}
+		if _, err := LoadRules([]*Rule{r1, r2}); err != nil {
+			t.Fatal(err)
+		}
+		rs2 := getRules()
+		if rs2[0].Resource == "abc1" {
+			assert.True(t, rs2[0] == r1)
+			assert.True(t, rs2[1] == r2)
+			assert.True(t, reflect.DeepEqual(rs2[0], r1))
+			assert.True(t, reflect.DeepEqual(rs2[1], r2))
+		} else {
+			assert.True(t, rs2[0] == r2)
+			assert.True(t, rs2[1] == r1)
+			assert.True(t, reflect.DeepEqual(rs2[0], r2))
+			assert.True(t, reflect.DeepEqual(rs2[1], r1))
+		}
+		if err := ClearRules(); err != nil {
+			t.Fatal(err)
+		}
+	})
 }