Add integration test for SentinelWebFluxFilter

Signed-off-by: Eric Zhao <sczyh16@gmail.com>

sentinel-adapter/sentinel-spring-webflux-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/webflux/SentinelWebFluxIntegrationTest.java

@@ -0,0 +1,169 @@
+package com.alibaba.csp.sentinel.adapter.spring.webflux;
+
+import java.util.ArrayList;
+import java.util.Collections;
+
+import com.alibaba.csp.sentinel.adapter.spring.webflux.callback.WebFluxCallbackManager;
+import com.alibaba.csp.sentinel.adapter.spring.webflux.test.WebFluxTestApplication;
+import com.alibaba.csp.sentinel.node.ClusterNode;
+import com.alibaba.csp.sentinel.slots.block.RuleConstant;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
+import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
+import com.alibaba.csp.sentinel.slots.clusterbuilder.ClusterBuilderSlot;
+import com.alibaba.csp.sentinel.util.StringUtil;
+
+import org.hamcrest.core.StringContains;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.SpringBootTest.WebEnvironment;
+import org.springframework.http.HttpStatus;
+import org.springframework.http.MediaType;
+import org.springframework.test.context.junit4.SpringRunner;
+import org.springframework.test.web.reactive.server.WebTestClient;
+import org.springframework.web.reactive.function.server.ServerResponse;
+
+import static org.junit.Assert.*;
+
+/**
+ * @author Eric Zhao
+ */
+@RunWith(SpringRunner.class)
+@SpringBootTest(classes = WebFluxTestApplication.class, webEnvironment = WebEnvironment.RANDOM_PORT)
+public class SentinelWebFluxIntegrationTest {
+
+    private static final String HELLO_STR = "Hello!";
+    private static final String BLOCK_MSG_PREFIX = "Blocked by Sentinel: ";
+
+    @Autowired
+    private WebTestClient webClient;
+
+    private void configureRulesFor(String resource, int count) {
+        configureRulesFor(resource, count, "default");
+    }
+
+    private void configureRulesFor(String resource, int count, String limitApp) {
+        FlowRule rule = new FlowRule()
+            .setCount(count)
+            .setGrade(RuleConstant.FLOW_GRADE_QPS);
+        rule.setResource(resource);
+        if (StringUtil.isNotBlank(limitApp)) {
+            rule.setLimitApp(limitApp);
+        }
+        FlowRuleManager.loadRules(Collections.singletonList(rule));
+    }
+
+    @Test
+    public void testWebFluxFilterBasic() throws Exception {
+        String url = "/hello";
+        this.webClient.get()
+            .uri(url)
+            .accept(MediaType.TEXT_PLAIN)
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).isEqualTo(HELLO_STR);
+
+        ClusterNode cn = ClusterBuilderSlot.getClusterNode(url);
+        assertNotNull(cn);
+        assertEquals(1, cn.passQps());
+    }
+
+    @Test
+    public void testCustomizedUrlCleaner() throws Exception {
+        final String fooPrefix = "/foo/";
+        String url1 = fooPrefix + 1;
+        String url2 = fooPrefix + 2;
+        WebFluxCallbackManager.setUrlCleaner(((exchange, originUrl) -> {
+            if (originUrl.startsWith(fooPrefix)) {
+                return "/foo/*";
+            }
+            return originUrl;
+        }));
+        this.webClient.get()
+            .uri(url1)
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).isEqualTo("Hello 1");
+        this.webClient.get()
+            .uri(url2)
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).isEqualTo("Hello 2");
+
+        ClusterNode cn = ClusterBuilderSlot.getClusterNode(fooPrefix + "*");
+        assertEquals(2, cn.passQps());
+        assertNull(ClusterBuilderSlot.getClusterNode(url1));
+        assertNull(ClusterBuilderSlot.getClusterNode(url2));
+
+        WebFluxCallbackManager.resetUrlCleaner();
+    }
+
+    @Test
+    public void testCustomizedBlockRequestHandler() throws Exception {
+        String url = "/error";
+        String prefix = "blocked: ";
+        WebFluxCallbackManager.setBlockHandler((exchange, t) -> ServerResponse.ok()
+            .contentType(MediaType.TEXT_PLAIN)
+            .syncBody(prefix + t.getMessage()));
+
+        this.webClient.get()
+            .uri(url)
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).value(StringContains.containsString(prefix));
+
+        WebFluxCallbackManager.resetBlockHandler();
+    }
+
+    @Test
+    public void testCustomizedRequestOriginParser() throws Exception {
+        String url = "/hello";
+        String limitOrigin = "userA";
+        final String headerName = "S-User";
+        configureRulesFor(url, 0, limitOrigin);
+
+        WebFluxCallbackManager.setRequestOriginParser(exchange -> {
+            String origin = exchange.getRequest().getHeaders().getFirst(headerName);
+            return origin != null ? origin : "";
+        });
+
+        this.webClient.get()
+            .uri(url)
+            .accept(MediaType.TEXT_PLAIN)
+            .header(headerName, "userB")
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).isEqualTo(HELLO_STR);
+        // This will be blocked.
+        this.webClient.get()
+            .uri(url)
+            .accept(MediaType.TEXT_PLAIN)
+            .header(headerName, limitOrigin)
+            .exchange()
+            .expectStatus().isEqualTo(HttpStatus.TOO_MANY_REQUESTS)
+            .expectBody(String.class).value(StringContains.containsString(BLOCK_MSG_PREFIX));
+        this.webClient.get()
+            .uri(url)
+            .accept(MediaType.TEXT_PLAIN)
+            .exchange()
+            .expectStatus().isOk()
+            .expectBody(String.class).isEqualTo(HELLO_STR);
+
+        WebFluxCallbackManager.resetRequestOriginParser();
+    }
+
+    @Before
+    public void setUp() {
+        FlowRuleManager.loadRules(new ArrayList<>());
+        ClusterBuilderSlot.resetClusterNodes();
+    }
+
+    @After
+    public void cleanUp() {
+        FlowRuleManager.loadRules(new ArrayList<>());
+        ClusterBuilderSlot.resetClusterNodes();
+    }
+}

sentinel-adapter/sentinel-spring-webflux-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/webflux/test/WebFluxTestApplication.java

@@ -0,0 +1,30 @@
+/*
+ * Copyright 1999-2019 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.adapter.spring.webflux.test;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+/**
+ * @author Eric Zhao
+ */
+@SpringBootApplication
+public class WebFluxTestApplication {
+
+    public static void main(String[] args) {
+        SpringApplication.run(WebFluxTestApplication.class, args);
+    }
+}

sentinel-adapter/sentinel-spring-webflux-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/webflux/test/WebFluxTestConfig.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 1999-2019 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.adapter.spring.webflux.test;
+
+import java.util.Collections;
+import java.util.List;
+
+import com.alibaba.csp.sentinel.adapter.spring.webflux.SentinelWebFluxFilter;
+import com.alibaba.csp.sentinel.adapter.spring.webflux.exception.SentinelBlockExceptionHandler;
+
+import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.annotation.Order;
+import org.springframework.http.codec.ServerCodecConfigurer;
+import org.springframework.web.reactive.result.view.ViewResolver;
+
+/**
+ * @author Eric Zhao
+ */
+@Configuration
+public class WebFluxTestConfig {
+
+    private final List<ViewResolver> viewResolvers;
+    private final ServerCodecConfigurer serverCodecConfigurer;
+
+    public WebFluxTestConfig(ObjectProvider<List<ViewResolver>> viewResolversProvider,
+                         ServerCodecConfigurer serverCodecConfigurer) {
+        this.viewResolvers = viewResolversProvider.getIfAvailable(Collections::emptyList);
+        this.serverCodecConfigurer = serverCodecConfigurer;
+    }
+
+    @Bean
+    @Order(-1)
+    public SentinelBlockExceptionHandler sentinelBlockExceptionHandler() {
+        // Register the block exception handler for Spring WebFlux.
+        return new SentinelBlockExceptionHandler(viewResolvers, serverCodecConfigurer);
+    }
+
+    @Bean
+    @Order(-1)
+    public SentinelWebFluxFilter sentinelWebFluxFilter() {
+        // Register the Sentinel WebFlux filter.
+        return new SentinelWebFluxFilter();
+    }
+}

sentinel-adapter/sentinel-spring-webflux-adapter/src/test/java/com/alibaba/csp/sentinel/adapter/spring/webflux/test/WebFluxTestController.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 1999-2019 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.adapter.spring.webflux.test;
+
+import com.alibaba.csp.sentinel.slots.block.flow.FlowException;
+
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.PathVariable;
+import org.springframework.web.bind.annotation.RestController;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+/**
+ * @author Eric Zhao
+ */
+@RestController
+public class WebFluxTestController {
+
+    @GetMapping("/hello")
+    public String apiHello() {
+        return "Hello!";
+    }
+
+    @GetMapping("/flux")
+    public Flux<Integer> apiFlux() {
+        return Flux.range(0, 5);
+    }
+
+    @GetMapping("/error")
+    public Mono<?> apiError() {
+        return Mono.error(new FlowException("testWebFluxError"));
+    }
+
+    @GetMapping("/foo/{id}")
+    public String apiFoo(@PathVariable("id") Long id) {
+        return "Hello " + id;
+    }
+}