Update test cases for flow slot and checker

Signed-off-by: Eric Zhao <[email protected]>

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleCheckerTest.java

@@ -0,0 +1,178 @@
+/*
+ * Copyright 1999-2018 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+import java.util.Arrays;
+
+import com.alibaba.csp.sentinel.EntryType;
+import com.alibaba.csp.sentinel.context.Context;
+import com.alibaba.csp.sentinel.node.ClusterNode;
+import com.alibaba.csp.sentinel.node.DefaultNode;
+import com.alibaba.csp.sentinel.slotchain.StringResourceWrapper;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.clusterbuilder.ClusterBuilderSlot;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+/**
+ * @author Eric Zhao
+ */
+public class FlowRuleCheckerTest {
+
+    @Test
+    public void testDefaultLimitAppFlowSelectNode() {
+        DefaultNode node = mock(DefaultNode.class);
+        ClusterNode cn = mock(ClusterNode.class);
+        when(node.getClusterNode()).thenReturn(cn);
+        Context context = mock(Context.class);
+
+        // limitApp: default
+        FlowRule rule = new FlowRule("testDefaultLimitAppFlowSelectNode").setCount(1);
+        assertEquals(cn, FlowRuleChecker.selectNodeByRequesterAndStrategy(rule, context, node));
+    }
+
+    @Test
+    public void testCustomOriginFlowSelectNode() {
+        String origin = "appA";
+        String limitAppB = "appB";
+
+        DefaultNode node = mock(DefaultNode.class);
+        DefaultNode originNode = mock(DefaultNode.class);
+        ClusterNode cn = mock(ClusterNode.class);
+        when(node.getClusterNode()).thenReturn(cn);
+        Context context = mock(Context.class);
+        when(context.getOrigin()).thenReturn(origin);
+        when(context.getOriginNode()).thenReturn(originNode);
+
+        FlowRule rule = new FlowRule("testCustomOriginFlowSelectNode").setCount(1);
+        rule.setLimitApp(origin);
+        // Origin matches, return the origin node.
+        assertEquals(originNode, FlowRuleChecker.selectNodeByRequesterAndStrategy(rule, context, node));
+
+        rule.setLimitApp(limitAppB);
+        // Origin mismatch, no node found.
+        assertNull(FlowRuleChecker.selectNodeByRequesterAndStrategy(rule, context, node));
+    }
+
+    @Test
+    public void testOtherOriginFlowSelectNode() {
+        String originA = "appA";
+        String originB = "appB";
+
+        DefaultNode node = mock(DefaultNode.class);
+        DefaultNode originNode = mock(DefaultNode.class);
+        ClusterNode cn = mock(ClusterNode.class);
+        when(node.getClusterNode()).thenReturn(cn);
+        Context context = mock(Context.class);
+        when(context.getOriginNode()).thenReturn(originNode);
+
+        FlowRule ruleA = new FlowRule("testOtherOriginFlowSelectNode").setCount(1);
+        ruleA.setLimitApp(originA);
+        FlowRule ruleB = new FlowRule("testOtherOriginFlowSelectNode").setCount(2);
+        ruleB.setLimitApp(RuleConstant.LIMIT_APP_OTHER);
+        FlowRuleManager.loadRules(Arrays.asList(ruleA, ruleB));
+
+        // Origin matches other, return the origin node.
+        when(context.getOrigin()).thenReturn(originB);
+        assertEquals(originNode, FlowRuleChecker.selectNodeByRequesterAndStrategy(ruleB, context, node));
+
+        // Origin matches limitApp of an existing rule, so no nodes are selected.
+        when(context.getOrigin()).thenReturn(originA);
+        assertNull(FlowRuleChecker.selectNodeByRequesterAndStrategy(ruleB, context, node));
+    }
+
+    @Test
+    public void testSelectNodeForEmptyReference() {
+        DefaultNode node = mock(DefaultNode.class);
+        Context context = mock(Context.class);
+
+        FlowRule rule = new FlowRule("testSelectNodeForEmptyReference")
+            .setCount(1)
+            .setStrategy(RuleConstant.STRATEGY_CHAIN);
+        assertNull(FlowRuleChecker.selectReferenceNode(rule, context, node));
+    }
+
+    @Test
+    public void testSelectNodeForRelateReference() {
+        String refResource = "testSelectNodeForRelateReference_refResource";
+
+        DefaultNode node = mock(DefaultNode.class);
+        ClusterNode refCn = mock(ClusterNode.class);
+        ClusterBuilderSlot.getClusterNodeMap().put(new StringResourceWrapper(refResource, EntryType.IN), refCn);
+        Context context = mock(Context.class);
+
+        FlowRule rule = new FlowRule("testSelectNodeForRelateReference")
+            .setCount(1)
+            .setStrategy(RuleConstant.STRATEGY_RELATE)
+            .setRefResource(refResource);
+        assertEquals(refCn, FlowRuleChecker.selectReferenceNode(rule, context, node));
+    }
+
+    @Test
+    public void testSelectReferenceNodeForContextEntrance() {
+        String contextName = "good_context";
+
+        DefaultNode node = mock(DefaultNode.class);
+        Context context = mock(Context.class);
+
+        FlowRule rule = new FlowRule("testSelectReferenceNodeForContextEntrance")
+            .setCount(1)
+            .setStrategy(RuleConstant.STRATEGY_CHAIN)
+            .setRefResource(contextName);
+
+        when(context.getName()).thenReturn(contextName);
+        assertEquals(node, FlowRuleChecker.selectReferenceNode(rule, context, node));
+
+        when(context.getName()).thenReturn("other_context");
+        assertNull(FlowRuleChecker.selectReferenceNode(rule, context, node));
+    }
+
+    @Test
+    public void testPassCheckNullLimitApp() {
+        FlowRule rule = new FlowRule("abc").setCount(1);
+        rule.setLimitApp(null);
+        assertTrue(FlowRuleChecker.passCheck(rule, null, null, 1));
+    }
+
+    @Test
+    public void testPassCheckSelectEmptyNodeSuccess() {
+        FlowRule rule = new FlowRule("abc").setCount(1);
+        rule.setLimitApp("abc");
+
+        DefaultNode node = mock(DefaultNode.class);
+        Context context = mock(Context.class);
+        when(context.getOrigin()).thenReturn("def");
+
+        assertTrue(FlowRuleChecker.passCheck(rule, context, node, 1));
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        FlowRuleManager.loadRules(null);
+        ClusterBuilderSlot.getClusterNodeMap().clear();
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        FlowRuleManager.loadRules(null);
+        ClusterBuilderSlot.getClusterNodeMap().clear();
+    }
+}

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowRuleTest.java

@@ -19,6 +19,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import org.junit.Ignore;
 import org.junit.Test;
 
 import com.alibaba.csp.sentinel.context.Context;
@@ -30,6 +31,7 @@
 /**
  * @author jialiang.linjl
  */
+@Ignore("Deprecated test for legacy FlowRule")
 public class FlowRuleTest {
 
     @Test

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/FlowSlotTest.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright 1999-2018 Alibaba Group Holding Ltd.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.alibaba.csp.sentinel.slots.block.flow;
+
+import java.util.Collections;
+
+import com.alibaba.csp.sentinel.EntryType;
+import com.alibaba.csp.sentinel.context.Context;
+import com.alibaba.csp.sentinel.context.ContextTestUtil;
+import com.alibaba.csp.sentinel.node.DefaultNode;
+import com.alibaba.csp.sentinel.slotchain.ResourceWrapper;
+import com.alibaba.csp.sentinel.slotchain.StringResourceWrapper;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+/**
+ * @author Eric Zhao
+ */
+public class FlowSlotTest {
+
+    @Before
+    public void setUp() {
+        ContextTestUtil.cleanUpContext();
+        FlowRuleManager.loadRules(null);
+    }
+
+    @After
+    public void tearDown() {
+        ContextTestUtil.cleanUpContext();
+        FlowRuleManager.loadRules(null);
+    }
+
+    @Test
+    public void testCheckFlowPass() throws Exception {
+        FlowSlot flowSlot = mock(FlowSlot.class);
+        Context context = mock(Context.class);
+        DefaultNode node = mock(DefaultNode.class);
+        doCallRealMethod().when(flowSlot).checkFlow(any(ResourceWrapper.class), any(Context.class),
+            any(DefaultNode.class), anyInt());
+
+        String resA = "resAK";
+        String resB = "resBK";
+        FlowRule rule1 = new FlowRule(resA).setCount(10);
+        FlowRule rule2 = new FlowRule(resB).setCount(10);
+        // Here we only load rules for resA.
+        FlowRuleManager.loadRules(Collections.singletonList(rule1));
+
+        when(flowSlot.canPassCheck(eq(rule1), any(Context.class), any(DefaultNode.class), anyInt()))
+            .thenReturn(true);
+        when(flowSlot.canPassCheck(eq(rule2), any(Context.class), any(DefaultNode.class), anyInt()))
+            .thenReturn(false);
+
+        flowSlot.checkFlow(new StringResourceWrapper(resA, EntryType.IN), context, node, 1);
+        flowSlot.checkFlow(new StringResourceWrapper(resB, EntryType.IN), context, node, 1);
+    }
+
+    @Test(expected = FlowException.class)
+    public void testCheckFlowBlock() throws Exception {
+        FlowSlot flowSlot = mock(FlowSlot.class);
+        Context context = mock(Context.class);
+        DefaultNode node = mock(DefaultNode.class);
+        doCallRealMethod().when(flowSlot).checkFlow(any(ResourceWrapper.class), any(Context.class),
+            any(DefaultNode.class), anyInt());
+
+        String resA = "resAK";
+        FlowRule rule = new FlowRule(resA).setCount(10);
+        FlowRuleManager.loadRules(Collections.singletonList(rule));
+
+        when(flowSlot.canPassCheck(any(FlowRule.class), any(Context.class), any(DefaultNode.class), anyInt()))
+            .thenReturn(false);
+
+        flowSlot.checkFlow(new StringResourceWrapper(resA, EntryType.IN), context, node, 1);
+    }
+}

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/DefaultControllerTest.java

@@ -0,0 +1,44 @@
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
+
+import com.alibaba.csp.sentinel.node.Node;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.TrafficShapingController;
+
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+import static org.mockito.Mockito.*;
+
+/**
+ * @author Eric Zhao
+ */
+public class DefaultControllerTest {
+
+    @Test
+    public void testCanPassForQps() {
+        long threshold = 10;
+        TrafficShapingController controller = new DefaultController(threshold, RuleConstant.FLOW_GRADE_QPS);
+        Node node = mock(Node.class);
+        when(node.passQps()).thenReturn(threshold - 1)
+            .thenReturn(threshold);
+
+        assertTrue(controller.canPass(node, 1));
+        assertFalse(controller.canPass(node, 1));
+    }
+
+    @Test
+    public void testCanPassForThreadCount() {
+        int threshold = 8;
+        TrafficShapingController controller = new DefaultController(threshold, RuleConstant.FLOW_GRADE_THREAD);
+        Node node = mock(Node.class);
+        when(node.curThreadNum()).thenReturn(threshold - 1)
+            .thenReturn(threshold);
+
+        assertTrue(controller.canPass(node, 1));
+        assertFalse(controller.canPass(node, 1));
+    }
+
+    @Test
+    public void testCanPassForQpsMultiThread() {
+    }
+}

sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/RateLimiterControllerTest.java → sentinel-core/src/test/java/com/alibaba/csp/sentinel/slots/block/flow/controller/RateLimiterControllerTest.java

@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.alibaba.csp.sentinel.slots.block.flow;
+package com.alibaba.csp.sentinel.slots.block.flow.controller;
 
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Mockito.mock;