@Test public void testSimpleCalls() throws Exception { FileSystem fs = Mockito.mock(FileSystem.class); Mockito.when(fs.getFileStatus(Mockito.any(Path.class))).thenReturn(new FileStatus(0, false, 0, 0, 0, new Path("/"))); Limiter limiter = new CountBasedLimiter(2); ThrottledFileSystem throttledFileSystem = new ThrottledFileSystem(fs, limiter, "testService"); Assert.assertNotNull(throttledFileSystem.getFileStatus(new Path("/myFile"))); Assert.assertNotNull(throttledFileSystem.getFileStatus(new Path("/myFile"))); try { throttledFileSystem.getFileStatus(new Path("/myFile")); Assert.fail(); } catch (NotEnoughPermitsException expected) { // Expected } }
@Override public FileSystem instrumentFileSystem(FileSystem fs, SharedResourcesBroker<S> broker, ConfigView<S, FileSystemKey> config) { try { String serviceName = ConfigUtils.getString(config.getConfig(), SERVICE_NAME_CONF_KEY, ""); Limiter limiter = broker.getSharedResource(new SharedLimiterFactory<S>(), new FileSystemLimiterKey(config.getKey().getUri())); return new ThrottledFileSystem(fs, limiter, serviceName); } catch (NotConfiguredException nce) { throw new RuntimeException(nce); } } }
@Test public void testListing() throws Exception { FileSystem fs = Mockito.mock(FileSystem.class); Mockito.when(fs.listStatus(Mockito.any(Path.class))).thenAnswer(new Answer<FileStatus[]>() { @Override public FileStatus[] answer(InvocationOnMock invocation) throws Throwable { Path path = (Path) invocation.getArguments()[0]; int files = Integer.parseInt(path.getName()); FileStatus status = new FileStatus(0, false, 0, 0, 0, new Path("/")); FileStatus[] out = new FileStatus[files]; for (int i = 0; i < files; i++) { out[i] = status; } return out; } }); Limiter limiter = new CountBasedLimiter(5); ThrottledFileSystem throttledFileSystem = new ThrottledFileSystem(fs, limiter, "testService"); Assert.assertEquals(throttledFileSystem.getServiceName(), "testService"); Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/99"))); // use 1 permit Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/250"))); // use 3 permits try { throttledFileSystem.listStatus(new Path("/files/150")); // requires 2 permits Assert.fail(); } catch (NotEnoughPermitsException expected) { // Expected } Assert.assertNotNull(throttledFileSystem.listStatus(new Path("/files/99"))); // requires 1 permit }
@Override public FileSystem instrumentFileSystem(FileSystem fs, SharedResourcesBroker<S> broker, ConfigView<S, FileSystemKey> config) { try { String serviceName = ConfigUtils.getString(config.getConfig(), SERVICE_NAME_CONF_KEY, ""); Limiter limiter = broker.getSharedResource(new SharedLimiterFactory<S>(), new FileSystemLimiterKey(config.getKey().getUri())); return new ThrottledFileSystem(fs, limiter, serviceName); } catch (NotConfiguredException nce) { throw new RuntimeException(nce); } } }