Search in sources :

Example 1 with MockClock

use of com.github.ambry.utils.MockClock in project ambry by linkedin.

the class HostLevelThrottlerTest method quotaTest.

/**
 * Test to make sure {@link HostLevelThrottler#shouldThrottle(RestRequest)} works as expected.
 */
@Test
public void quotaTest() throws Exception {
    Properties props = new Properties();
    props.setProperty(HostThrottleConfig.REST_REQUEST_QUOTA_STRING, "{\"PUT\": \"20\",\"GET\": \"20\",\"POST\": \"20\",\"HEAD\": \"20\",\"OPTIONS\": \"20\",\"DELETE\": \"20\"}");
    HostThrottleConfig hostThrottleConfig = new HostThrottleConfig(new VerifiableProperties(props));
    MockClock clock = new MockClock();
    HostLevelThrottler quotaManager = new HostLevelThrottler(createQuotaMock(hostThrottleConfig, clock), new HashMap<>(), null);
    // Issue new requests. Since MockClock tick doesn't change, rate is 0.
    for (int i = 0; i < 100; i++) {
        for (RestMethod restMethod : RestMethod.values()) {
            RestRequest restRequest = createRestRequest(restMethod, "http://www.linkedin.com/");
            Assert.assertFalse("Should not throttle", quotaManager.shouldThrottle(restRequest));
        }
    }
    // Move MockClock ahead to 5 seconds later. Rate = 20. New requests should be denied unless its quota is not defined.
    clock.tick(5);
    for (RestMethod restMethod : RestMethod.values()) {
        RestRequest restRequest = createRestRequest(restMethod, "http://www.linkedin.com/");
        if (restMethod == RestMethod.UNKNOWN) {
            Assert.assertFalse("Should not throttle.", quotaManager.shouldThrottle(restRequest));
        } else {
            Assert.assertTrue("Should throttle", quotaManager.shouldThrottle(restRequest));
        }
    }
    // Clock tick to another 5 seconds later, rate < 20. Accept new requests.
    clock.tick(5);
    for (RestMethod restMethod : RestMethod.values()) {
        RestRequest restRequest = createRestRequest(restMethod, "http://www.linkedin.com/");
        Assert.assertFalse("Should not throttle", quotaManager.shouldThrottle(restRequest));
    }
}
Also used : MockRestRequest(com.github.ambry.rest.MockRestRequest) RestRequest(com.github.ambry.rest.RestRequest) VerifiableProperties(com.github.ambry.config.VerifiableProperties) Properties(java.util.Properties) VerifiableProperties(com.github.ambry.config.VerifiableProperties) HostThrottleConfig(com.github.ambry.config.HostThrottleConfig) MockClock(com.github.ambry.utils.MockClock) RestMethod(com.github.ambry.rest.RestMethod) Test(org.junit.Test)

Aggregations

HostThrottleConfig (com.github.ambry.config.HostThrottleConfig)1 VerifiableProperties (com.github.ambry.config.VerifiableProperties)1 MockRestRequest (com.github.ambry.rest.MockRestRequest)1 RestMethod (com.github.ambry.rest.RestMethod)1 RestRequest (com.github.ambry.rest.RestRequest)1 MockClock (com.github.ambry.utils.MockClock)1 Properties (java.util.Properties)1 Test (org.junit.Test)1