Examples with RestRequest com.github.ambry.rest.RestRequest used on opensource projects

Example 96 with RestRequest

use of com.github.ambry.rest.RestRequest in project ambry by linkedin.

the class TtlUpdateHandlerTest method handleGoodCaseTest.

/**
 * Tests the case where TTL update succeeds
 * @throws Exception
 */
@Test
public void handleGoodCaseTest() throws Exception {
    RestRequest restRequest = new MockRestRequest(MockRestRequest.DUMMY_DATA, null);
    restRequest.setArg(RestUtils.Headers.BLOB_ID, blobId);
    restRequest.setArg(RestUtils.Headers.SERVICE_ID, SERVICE_ID);
    verifyTtlUpdate(restRequest, REF_ACCOUNT, REF_CONTAINER);
}

Example 97 with RestRequest

use of com.github.ambry.rest.RestRequest in project ambry by linkedin.

the class TtlUpdateHandlerTest method verifyFailureWithMsg.

/**
 * Verifies that attempting to update TTL fails with the provided {@code msg}.
 * @param msg the message in the {@link Exception} that will be thrown.
 * @throws Exception
 */
private void verifyFailureWithMsg(String msg) throws Exception {
    RestRequest restRequest = new MockRestRequest(MockRestRequest.DUMMY_DATA, null);
    restRequest.setArg(RestUtils.Headers.BLOB_ID, blobId);
    restRequest.setArg(RestUtils.Headers.SERVICE_ID, SERVICE_ID);
    try {
        sendRequestGetResponse(restRequest, new MockRestResponseChannel());
        fail("Request should have failed");
    } catch (Exception e) {
        if (!msg.equals(e.getMessage())) {
            throw e;
        }
    }
}

Example 98 with RestRequest

use of com.github.ambry.rest.RestRequest in project ambry by linkedin.

the class UndeleteHandlerTest method verifyFailureWithMsg.

/**
 * Verifies that attempting to undelete fails with the provided {@code msg}.
 * @param msg the message in the {@link Exception} that will be thrown.
 * @throws Exception
 */
private void verifyFailureWithMsg(String msg) throws Exception {
    RestRequest restRequest = new MockRestRequest(MockRestRequest.DUMMY_DATA, null);
    restRequest.setArg(RestUtils.Headers.BLOB_ID, blobId);
    restRequest.setArg(RestUtils.Headers.SERVICE_ID, SERVICE_ID);
    try {
        sendRequestGetResponse(restRequest, new MockRestResponseChannel());
        fail("Request should have failed");
    } catch (Exception e) {
        if (!msg.equals(e.getMessage())) {
            throw e;
        }
    }
}

Example 99 with RestRequest

use of com.github.ambry.rest.RestRequest in project ambry by linkedin.

the class StorageQuotaEnforcerTest method testGetQuotaAndUsageAndCharge.

/**
 * Test {@link StorageQuotaEnforcer#getQuotaAndUsage} and {@link StorageQuotaEnforcer#charge} methods.
 * @throws Exception
 */
@Test
public void testGetQuotaAndUsageAndCharge() throws Exception {
    int initNumAccounts = 10;
    Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(initNumAccounts, 10, 10000, 1000);
    InMemAccountService accountService = new InMemAccountService(false, false);
    Map<String, Long> storageQuota = new HashMap<>();
    // Account and container id's base is 1, not 0
    for (int i = 1; i <= initNumAccounts; i++) {
        QuotaResourceType resourceType = i <= containerUsage.size() / 2 ? QuotaResourceType.CONTAINER : QuotaResourceType.ACCOUNT;
        AccountBuilder accountBuilder = new AccountBuilder((short) i, String.valueOf(i), Account.AccountStatus.ACTIVE, resourceType);
        for (int j = 1; j <= 10; j++) {
            accountBuilder.addOrUpdateContainer(new ContainerBuilder((short) j, String.valueOf(j), Container.ContainerStatus.ACTIVE, "", (short) i).build());
        }
        accountService.updateAccounts(Collections.singleton(accountBuilder.build()));
        if (resourceType == QuotaResourceType.ACCOUNT) {
            storageQuota.put(QuotaResource.fromAccount(accountService.getAccountById((short) i)).getResourceId(), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
        } else {
            accountService.getAccountById((short) i).getAllContainers().forEach(c -> storageQuota.put(QuotaResource.fromContainer(c).getResourceId(), containerUsage.get(String.valueOf(c.getParentAccountId())).get(String.valueOf(c.getId()))));
        }
    }
    JSONStringStorageQuotaSource quotaSource = new JSONStringStorageQuotaSource(storageQuota, accountService);
    StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, quotaSource, (StorageUsageRefresher) null);
    enforcer.initStorageUsage(Collections.EMPTY_MAP);
    for (Map.Entry<String, Map<String, Long>> accountEntry : containerUsage.entrySet()) {
        short accountId = Short.valueOf(accountEntry.getKey());
        if (accountService.getAccountById(accountId).getQuotaResourceType() == QuotaResourceType.ACCOUNT) {
            long quota = (long) (quotaSource.getQuota(QuotaResource.fromAccount(accountService.getAccountById(accountId)), QuotaName.STORAGE_IN_GB).getQuotaValue()) * BYTES_IN_GB;
            RestRequest restRequest = createRestRequest(accountService, accountId, (short) 1);
            Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
            assertEquals("Account id: " + accountEntry.getKey(), quota, quotaAndUsage.getFirst().longValue());
            assertEquals(0L, quotaAndUsage.getSecond().longValue());
            quotaAndUsage = enforcer.charge(restRequest, quota / 2);
            assertEquals(quota, quotaAndUsage.getFirst().longValue());
            assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
            quotaAndUsage = enforcer.charge(restRequest, quota);
            assertEquals(quota, quotaAndUsage.getFirst().longValue());
            assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
        } else {
            for (Map.Entry<String, Long> containerEntry : accountEntry.getValue().entrySet()) {
                short containerId = Short.valueOf(containerEntry.getKey());
                long quota = containerEntry.getValue() * BYTES_IN_GB;
                RestRequest restRequest = createRestRequest(accountService, accountId, containerId);
                Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(0L, quotaAndUsage.getSecond().longValue());
                quotaAndUsage = enforcer.charge(restRequest, quota / 2);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
                quotaAndUsage = enforcer.charge(restRequest, quota);
                assertEquals(quota, quotaAndUsage.getFirst().longValue());
                assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
            }
        }
    }
    // Now create a restRequest that doesn't carry account and container
    RestRequest restRequest = createRestRequest();
    Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    quotaAndUsage = enforcer.charge(restRequest, 100L);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    Account account = new AccountBuilder((short) 1000, String.valueOf(1000), Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 10000, String.valueOf(10000), Container.ContainerStatus.ACTIVE, "", (short) 1000).build()).build();
    accountService.updateAccounts(Collections.singleton(account));
    restRequest = createRestRequest(accountService, (short) 1000, (short) 10000);
    quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
    quotaAndUsage = enforcer.charge(restRequest, 100L);
    assertEquals(-1L, quotaAndUsage.getFirst().longValue());
    assertEquals(0L, quotaAndUsage.getSecond().longValue());
}

Example 100 with RestRequest

use of com.github.ambry.rest.RestRequest in project ambry by linkedin.

the class AmbryQuotaManagerTest method testRecommend.

@Test
public void testRecommend() throws Exception {
    RestRequest restRequest = QuotaTestUtils.createRestRequest(ACCOUNT, ACCOUNT.getAllContainers().iterator().next(), RestMethod.GET);
    AmbryQuotaManager ambryQuotaManager = buildAmbryQuotaManagerWithNoEnforcers();
    ambryQuotaManager.init();
    // 1. test that recommend returns null with not enforcers.
    Assert.assertNull(ambryQuotaManager.recommend(restRequest));
    // 2. test that enforcers are called.
    ambryQuotaManager = buildAmbryQuotaManagerWithAmbryCUEnforcer();
    testAmbryCUQuotaEnforcer.resetDefaults();
    ambryQuotaManager.init();
    ThrottlingRecommendation throttlingRecommendation = ambryQuotaManager.recommend(restRequest);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer.initCallCount);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer.recommendCallCount);
    Assert.assertEquals(throttlingRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    Assert.assertEquals(2, TestCUQuotaEnforcerFactory.TestQuotaRecommendationMergePolicy.quotaRecommendationsCount);
    // 3. test that if any quota enforcer throws exception, then other quota enforcers are called nonetheless
    testAmbryCUQuotaEnforcer.resetDefaults();
    testAmbryCUQuotaEnforcer.throwException = true;
    throttlingRecommendation = ambryQuotaManager.recommend(restRequest);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer.recommendCallCount);
    Assert.assertEquals(throttlingRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    Assert.assertEquals(3, TestCUQuotaEnforcerFactory.TestQuotaRecommendationMergePolicy.quotaRecommendationsCount);
    ambryQuotaManager = buildAmbryQuotaManager();
    List<QuotaEnforcer> quotaEnforcers = new ArrayList<>(ambryQuotaManager.quotaEnforcers);
    TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer testAmbryCUQuotaEnforcer1 = (TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer) quotaEnforcers.get(0);
    TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer testAmbryCUQuotaEnforcer2 = (TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer) quotaEnforcers.get(1);
    testAmbryCUQuotaEnforcer1.resetDefaults();
    testAmbryCUQuotaEnforcer2.resetDefaults();
    ambryQuotaManager.init();
    // 4. test that if any quota enforcer recommends DELAY action, then that's the recommendation
    testAmbryCUQuotaEnforcer1.resetDefaults();
    testAmbryCUQuotaEnforcer2.resetDefaults();
    testAmbryCUQuotaEnforcer1.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.DELAY_QUOTA_RECOMMENDATION;
    testAmbryCUQuotaEnforcer2.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.ALLOW_QUOTA_RECOMMENDATION;
    throttlingRecommendation = ambryQuotaManager.recommend(restRequest);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
    Assert.assertEquals(throttlingRecommendation.getQuotaAction(), QuotaAction.DELAY);
    Assert.assertEquals(5, TestCUQuotaEnforcerFactory.TestQuotaRecommendationMergePolicy.quotaRecommendationsCount);
    // 5. test that if any quota enforcer recommends REJECT action, then that's the recommendation
    testAmbryCUQuotaEnforcer1.resetDefaults();
    testAmbryCUQuotaEnforcer2.resetDefaults();
    testAmbryCUQuotaEnforcer1.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.DELAY_QUOTA_RECOMMENDATION;
    testAmbryCUQuotaEnforcer2.recommendReturnVal = TestCUQuotaEnforcerFactory.TestAmbryCUQuotaEnforcer.REJECT_QUOTA_RECOMMENDATION;
    throttlingRecommendation = ambryQuotaManager.recommend(restRequest);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
    Assert.assertEquals(throttlingRecommendation.getQuotaAction(), QuotaAction.REJECT);
    Assert.assertEquals(7, TestCUQuotaEnforcerFactory.TestQuotaRecommendationMergePolicy.quotaRecommendationsCount);
    // 6. test if all the enforcers throw exception.
    testAmbryCUQuotaEnforcer1.resetDefaults();
    testAmbryCUQuotaEnforcer2.resetDefaults();
    testAmbryCUQuotaEnforcer1.throwException = true;
    testAmbryCUQuotaEnforcer2.throwException = true;
    throttlingRecommendation = ambryQuotaManager.recommend(restRequest);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer1.recommendCallCount);
    Assert.assertEquals(1, testAmbryCUQuotaEnforcer2.recommendCallCount);
    Assert.assertEquals(throttlingRecommendation.getQuotaAction(), QuotaAction.ALLOW);
    Assert.assertEquals(7, TestCUQuotaEnforcerFactory.TestQuotaRecommendationMergePolicy.quotaRecommendationsCount);
}