use of org.apache.gobblin.restli.throttling.ThrottlingPolicyFactory 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();
}
}
}
Aggregations