Search in sources :

Example 1 with EmbeddedRestliServer

use of org.apache.gobblin.restli.EmbeddedRestliServer in project incubator-gobblin by apache.

the class ThrottlingClientTest method test.

@Test
public void test() throws Exception {
    ThrottlingPolicyFactory factory = new ThrottlingPolicyFactory();
    SharedLimiterKey res1key = new SharedLimiterKey("res1");
    Map<String, String> configMap = ImmutableMap.<String, String>builder().put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null, ThrottlingPolicyFactory.POLICY_KEY), CountBasedPolicy.FACTORY_ALIAS).put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null, CountBasedPolicy.COUNT_KEY), "50").put(BrokerConfigurationKeyGenerator.generateKey(factory, null, null, ThrottlingPolicyFactory.FAIL_ON_UNKNOWN_RESOURCE_ID), "true").build();
    ThrottlingGuiceServletConfig guiceServletConfig = new ThrottlingGuiceServletConfig();
    guiceServletConfig.initialize(ConfigFactory.parseMap(configMap));
    Injector injector = guiceServletConfig.getInjector();
    EmbeddedRestliServer server = EmbeddedRestliServer.builder().resources(Lists.<Class<? extends BaseResource>>newArrayList(LimiterServerResource.class)).injector(injector).build();
    try {
        server.startAsync();
        server.awaitRunning();
        final HttpClientFactory http = new HttpClientFactory();
        final Client r2Client = new TransportClientAdapter(http.getClient(Collections.<String, String>emptyMap()));
        RestClient restClient = new RestClient(r2Client, server.getURIPrefix());
        PermitsGetRequestBuilder getBuilder = new PermitsRequestBuilders().get();
        PermitRequest res1request = new PermitRequest();
        res1request.setPermits(20);
        res1request.setResource(res1key.getResourceLimitedPath());
        PermitAllocation allocation = getPermitAllocation(res1request, restClient, getBuilder);
        Assert.assertEquals(allocation.getPermits(), new Long(20));
        allocation = getPermitAllocation(res1request, restClient, getBuilder);
        Assert.assertEquals(allocation.getPermits(), new Long(20));
        // out of permits
        try {
            allocation = getPermitAllocation(res1request, restClient, getBuilder);
            Assert.fail();
        } catch (RestLiResponseException exc) {
            Assert.assertEquals(exc.getStatus(), HttpStatus.S_403_FORBIDDEN.getCode());
        }
        PermitRequest invalidRequest = new PermitRequest();
        invalidRequest.setPermits(20);
        invalidRequest.setResource("invalidkey");
        try {
            allocation = getPermitAllocation(invalidRequest, restClient, getBuilder);
            Assert.fail();
        } catch (RestLiResponseException exc) {
            Assert.assertEquals(exc.getStatus(), 422);
        }
    } finally {
        if (server.isRunning()) {
            server.stopAsync();
            server.awaitTerminated();
        }
    }
}
Also used : EmbeddedRestliServer(org.apache.gobblin.restli.EmbeddedRestliServer) RestClient(com.linkedin.restli.client.RestClient) Injector(com.google.inject.Injector) TransportClientAdapter(com.linkedin.r2.transport.common.bridge.client.TransportClientAdapter) RestLiResponseException(com.linkedin.restli.client.RestLiResponseException) SharedLimiterKey(org.apache.gobblin.util.limiter.broker.SharedLimiterKey) RestClient(com.linkedin.restli.client.RestClient) Client(com.linkedin.r2.transport.common.Client) HttpClientFactory(com.linkedin.r2.transport.http.client.HttpClientFactory) Test(org.testng.annotations.Test)

Example 2 with EmbeddedRestliServer

use of org.apache.gobblin.restli.EmbeddedRestliServer in project incubator-gobblin by apache.

the class RestliServiceBasedLimiterTest method test.

@Test
public void test() throws Exception {
    ThrottlingPolicyFactory factory = new ThrottlingPolicyFactory();
    SharedLimiterKey res1key = new SharedLimiterKey("res1");
    Map<String, String> configMap = ImmutableMap.<String, String>builder().put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null, ThrottlingPolicyFactory.POLICY_KEY), CountBasedPolicy.FACTORY_ALIAS).put(BrokerConfigurationKeyGenerator.generateKey(factory, res1key, null, CountBasedPolicy.COUNT_KEY), "100").build();
    ThrottlingGuiceServletConfig guiceServletConfig = new ThrottlingGuiceServletConfig();
    guiceServletConfig.initialize(ConfigFactory.parseMap(configMap));
    Injector injector = guiceServletConfig.getInjector();
    EmbeddedRestliServer server = EmbeddedRestliServer.builder().resources(Lists.<Class<? extends BaseResource>>newArrayList(LimiterServerResource.class)).injector(injector).build();
    SharedResourcesBroker<SimpleScopeType> broker = SharedResourcesBrokerFactory.createDefaultTopLevelBroker(ConfigFactory.empty(), SimpleScopeType.GLOBAL.defaultScopeInstance());
    try {
        server.startAsync();
        server.awaitRunning();
        RestliServiceBasedLimiter limiter = RestliServiceBasedLimiter.builder().requestSender(new RedirectAwareRestClientRequestSender(broker, Lists.newArrayList(server.getURIPrefix()))).resourceLimited(res1key.getResourceLimitedPath()).serviceIdentifier("service").build();
        Assert.assertNotNull(limiter.acquirePermits(20));
        Assert.assertNotNull(limiter.acquirePermits(20));
        Assert.assertNull(limiter.acquirePermits(1000));
    } finally {
        if (server.isRunning()) {
            server.stopAsync();
            server.awaitTerminated();
        }
    }
}
Also used : ThrottlingGuiceServletConfig(org.apache.gobblin.restli.throttling.ThrottlingGuiceServletConfig) ThrottlingPolicyFactory(org.apache.gobblin.restli.throttling.ThrottlingPolicyFactory) EmbeddedRestliServer(org.apache.gobblin.restli.EmbeddedRestliServer) Injector(com.google.inject.Injector) LimiterServerResource(org.apache.gobblin.restli.throttling.LimiterServerResource) SharedLimiterKey(org.apache.gobblin.util.limiter.broker.SharedLimiterKey) SimpleScopeType(org.apache.gobblin.broker.SimpleScopeType) Test(org.testng.annotations.Test)

Example 3 with EmbeddedRestliServer

use of org.apache.gobblin.restli.EmbeddedRestliServer in project incubator-gobblin by apache.

the class RestliServiceBasedLimiterTest method createAndStartServer.

private RestliServer createAndStartServer(Config baseConfig, int port) {
    ThrottlingGuiceServletConfig guiceServletConfig = new ThrottlingGuiceServletConfig();
    guiceServletConfig.initialize(baseConfig.withFallback(ConfigFactory.parseMap(ImmutableMap.of(ThrottlingGuiceServletConfig.LISTENING_PORT, Integer.toString(port)))));
    Injector injector = guiceServletConfig.getInjector();
    EmbeddedRestliServer server = EmbeddedRestliServer.builder().resources(Lists.<Class<? extends BaseResource>>newArrayList(LimiterServerResource.class)).injector(injector).port(port).build();
    server.startAsync();
    server.awaitRunning();
    return new RestliServer(server, guiceServletConfig);
}
Also used : ThrottlingGuiceServletConfig(org.apache.gobblin.restli.throttling.ThrottlingGuiceServletConfig) EmbeddedRestliServer(org.apache.gobblin.restli.EmbeddedRestliServer) EmbeddedRestliServer(org.apache.gobblin.restli.EmbeddedRestliServer) Injector(com.google.inject.Injector) BaseResource(com.linkedin.restli.server.resources.BaseResource)

Aggregations

Injector (com.google.inject.Injector)3 EmbeddedRestliServer (org.apache.gobblin.restli.EmbeddedRestliServer)3 ThrottlingGuiceServletConfig (org.apache.gobblin.restli.throttling.ThrottlingGuiceServletConfig)2 SharedLimiterKey (org.apache.gobblin.util.limiter.broker.SharedLimiterKey)2 Test (org.testng.annotations.Test)2 Client (com.linkedin.r2.transport.common.Client)1 TransportClientAdapter (com.linkedin.r2.transport.common.bridge.client.TransportClientAdapter)1 HttpClientFactory (com.linkedin.r2.transport.http.client.HttpClientFactory)1 RestClient (com.linkedin.restli.client.RestClient)1 RestLiResponseException (com.linkedin.restli.client.RestLiResponseException)1 BaseResource (com.linkedin.restli.server.resources.BaseResource)1 SimpleScopeType (org.apache.gobblin.broker.SimpleScopeType)1 LimiterServerResource (org.apache.gobblin.restli.throttling.LimiterServerResource)1 ThrottlingPolicyFactory (org.apache.gobblin.restli.throttling.ThrottlingPolicyFactory)1