Examples with QuotaRequest com.tencent.polaris.ratelimit.api.rpc.QuotaRequest used on opensource projects

Example 6 with QuotaRequest

use of com.tencent.polaris.ratelimit.api.rpc.QuotaRequest in project spring-cloud-tencent by Tencent.

the class QuotaCheckServletFilterTest method setUp.

@Before
public void setUp() {
    MetadataContext.LOCAL_NAMESPACE = "TEST";
    LimitAPI limitAPI = mock(LimitAPI.class);
    when(limitAPI.getQuota(any(QuotaRequest.class))).thenAnswer(invocationOnMock -> {
        String serviceName = ((QuotaRequest) invocationOnMock.getArgument(0)).getService();
        if (serviceName.equals("TestApp1")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 0, "QuotaResultOk"));
        } else if (serviceName.equals("TestApp2")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultOk, 1000, "QuotaResultOk"));
        } else if (serviceName.equals("TestApp3")) {
            return new QuotaResponse(new QuotaResult(QuotaResult.Code.QuotaResultLimited, 0, "QuotaResultLimited"));
        } else {
            return new QuotaResponse(new QuotaResult(null, 0, null));
        }
    });
    PolarisRateLimitProperties polarisRateLimitProperties = new PolarisRateLimitProperties();
    polarisRateLimitProperties.setRejectRequestTips("RejectRequestTips");
    polarisRateLimitProperties.setRejectHttpCode(419);
    RateLimitRuleLabelResolver rateLimitRuleLabelResolver = mock(RateLimitRuleLabelResolver.class);
    when(rateLimitRuleLabelResolver.getExpressionLabelKeys(anyString(), anyString())).thenReturn(Collections.EMPTY_SET);
    this.quotaCheckServletFilter = new QuotaCheckServletFilter(limitAPI, labelResolver, polarisRateLimitProperties, rateLimitRuleLabelResolver);
}

Example 7 with QuotaRequest

use of com.tencent.polaris.ratelimit.api.rpc.QuotaRequest in project polaris-java by polarismesh.

the class APIFacade method getQuota.

public static boolean getQuota(String namespace, String service, String method, Map<String, String> labels, int count) {
    if (!inited.get()) {
        LOGGER.info("polaris not inited, getQuota fail");
        return false;
    }
    QuotaRequest quotaRequest = new QuotaRequest();
    quotaRequest.setNamespace(namespace);
    quotaRequest.setService(service);
    quotaRequest.setMethod(method);
    quotaRequest.setLabels(labels);
    quotaRequest.setCount(count);
    QuotaResponse quota = limitAPI.getQuota(quotaRequest);
    return quota.getCode() == QuotaResultCode.QuotaResultOk;
}

Example 8 with QuotaRequest

use of com.tencent.polaris.ratelimit.api.rpc.QuotaRequest in project polaris-java by polarismesh.

the class LocalTest method testQuotaAcquire.

private static void testQuotaAcquire(LimitAPI limitAPI, Map<String, String> labels, int maxCount) {
    QuotaRequest payRequest = new QuotaRequest();
    payRequest.setNamespace(Consts.NAMESPACE_TEST);
    payRequest.setService(Consts.LOCAL_LIMIT_SERVICE);
    payRequest.setCount(1);
    payRequest.setLabels(labels);
    boolean payLimit = false;
    boolean payPass = false;
    for (int i = 0; i < maxCount; i++) {
        QuotaResponse response = limitAPI.getQuota(payRequest);
        if (response.getCode() == QuotaResultCode.QuotaResultOk) {
            payPass = true;
        } else if (response.getCode() == QuotaResultCode.QuotaResultLimited) {
            payLimit = true;
        }
    }
    Assert.assertTrue(payPass);
    Assert.assertTrue(payLimit);
}