pull/902/head
seanyu 3 years ago
parent ec5fc62a9d
commit a8d54a0e9b

@ -17,6 +17,8 @@
package com.tencent.cloud.metadata.core;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Map;
import com.tencent.cloud.common.metadata.MetadataContext;
@ -27,9 +29,12 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.http.HttpMethod;
import org.springframework.mock.http.client.MockClientHttpRequest;
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import org.springframework.web.reactive.function.client.ClientRequest;
import static org.assertj.core.api.Assertions.assertThat;
import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.RANDOM_PORT;
@ -82,4 +87,23 @@ public class TransHeadersTransferTest {
assertThat(transHeadersKV.get("header2")).isEqualTo(JacksonUtils.serialize2Json(header2));
assertThat(transHeadersKV.get("header3")).isEqualTo(JacksonUtils.serialize2Json(header3));
}
@Test
public void transferWebClientTest() throws URISyntaxException {
MetadataContext metadataContext = MetadataContextHolder.get();
metadataContext.setTransHeaders("header1,header2,header3", "");
ClientRequest.Builder builder = ClientRequest.create(HttpMethod.GET, new URI(""));
String[] header1 = {"1"};
String[] header2 = {"2"};
String[] header3 = {"3"};
builder.header("header1", header1);
builder.header("header2", header2);
builder.header("header3", header3);
ClientRequest request = builder.build();
TransHeadersTransfer.transfer(request);
Map<String, String> transHeadersKV = MetadataContextHolder.get().getTransHeadersKV();
assertThat(transHeadersKV.get("header1")).isEqualTo(JacksonUtils.serialize2Json(header1));
assertThat(transHeadersKV.get("header2")).isEqualTo(JacksonUtils.serialize2Json(header2));
assertThat(transHeadersKV.get("header3")).isEqualTo(JacksonUtils.serialize2Json(header3));
}
}

@ -23,6 +23,7 @@ import java.io.InputStreamReader;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
@ -30,6 +31,7 @@ import java.util.stream.Collectors;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelReactiveResolver;
@ -46,6 +48,8 @@ import org.springframework.mock.web.server.MockServerWebExchange;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import org.springframework.web.server.ServerWebExchange;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
@ -92,6 +96,12 @@ public class RateLimitRuleArgumentReactiveResolverTest {
.queryParam("yyy", "yyy")
.build();
ServerWebExchange exchange = MockServerWebExchange.from(request);
MetadataContext metadataContext = new MetadataContext();
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
}});
MetadataContextHolder.set(metadataContext);
Set<Argument> arguments = rateLimitRuleArgumentReactiveResolver.getArguments(exchange, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
Set<Argument> exceptRes = new HashSet<>();
exceptRes.add(Argument.buildMethod("GET"));
@ -99,6 +109,7 @@ public class RateLimitRuleArgumentReactiveResolverTest {
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
assertThat(arguments).isEqualTo(exceptRes);
}

@ -22,6 +22,7 @@ import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.stream.Collectors;
@ -29,6 +30,7 @@ import java.util.stream.Collectors;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import com.tencent.cloud.common.metadata.MetadataContext;
import com.tencent.cloud.common.metadata.MetadataContextHolder;
import com.tencent.cloud.polaris.context.ServiceRuleManager;
import com.tencent.cloud.polaris.ratelimit.filter.QuotaCheckServletFilterTest;
import com.tencent.cloud.polaris.ratelimit.spi.PolarisRateLimiterLabelServletResolver;
@ -43,6 +45,8 @@ import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.mock.web.MockHttpServletRequest;
import org.springframework.test.context.junit.jupiter.SpringExtension;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAME;
import static com.tencent.cloud.common.constant.MetadataConstant.DefaultMetadata.DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
@ -84,6 +88,12 @@ public class RateLimitRuleArgumentServletResolverTest {
MockHttpServletRequest request = new MockHttpServletRequest(null, "GET", "/xxx");
request.setParameter("yyy", "yyy");
request.addHeader("xxx", "xxx");
MetadataContext metadataContext = new MetadataContext();
metadataContext.setUpstreamDisposableMetadata(new HashMap<String, String>() {{
put(DEFAULT_METADATA_SOURCE_SERVICE_NAMESPACE, MetadataContext.LOCAL_NAMESPACE);
put(DEFAULT_METADATA_SOURCE_SERVICE_NAME, MetadataContext.LOCAL_SERVICE);
}});
MetadataContextHolder.set(metadataContext);
Set<Argument> arguments = rateLimitRuleArgumentServletResolver.getArguments(request, MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE);
Set<Argument> exceptRes = new HashSet<>();
exceptRes.add(Argument.buildMethod("GET"));
@ -91,6 +101,7 @@ public class RateLimitRuleArgumentServletResolverTest {
exceptRes.add(Argument.buildQuery("yyy", "yyy"));
exceptRes.add(Argument.buildCallerIP("127.0.0.1"));
exceptRes.add(Argument.buildCustom("xxx", "xxx"));
exceptRes.add(Argument.buildCallerService(MetadataContext.LOCAL_NAMESPACE, MetadataContext.LOCAL_SERVICE));
assertThat(arguments).isEqualTo(exceptRes);
}

@ -17,6 +17,7 @@
package com.tencent.cloud.polaris.ratelimit.utils;
import java.util.HashMap;
import java.util.HashSet;
import com.tencent.polaris.api.plugin.ratelimiter.QuotaResult;
@ -66,6 +67,37 @@ public class QuotaCheckUtilsTest {
@Test
public void testGetQuota() {
// Pass
String serviceName = "TestApp1";
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashMap<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultOk");
// Unirate waiting 1000ms
serviceName = "TestApp2";
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashMap<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
assertThat(quotaResponse.getWaitMs()).isEqualTo(1000);
assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultOk");
// Rate limited
serviceName = "TestApp3";
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashMap<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultLimited);
assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
assertThat(quotaResponse.getInfo()).isEqualTo("QuotaResultLimited");
// Exception
serviceName = "TestApp4";
quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashMap<>(), null);
assertThat(quotaResponse.getCode()).isEqualTo(QuotaResultCode.QuotaResultOk);
assertThat(quotaResponse.getWaitMs()).isEqualTo(0);
assertThat(quotaResponse.getInfo()).isEqualTo("get quota failed");
}
@Test
public void testGetQuota2() {
// Pass
String serviceName = "TestApp1";
QuotaResponse quotaResponse = QuotaCheckUtils.getQuota(limitAPI, null, serviceName, 1, new HashSet<>(), null);

Loading…
Cancel
Save