QuotaService extends CompositeService { public QuotaService(RMContext rmContext) { super(QuotaService.class.getName()); this.rmContext = rmContext; } QuotaService(RMContext rmContext); @SuppressWarnings("unchecked") void containerStarted(RMContainer container); @SuppressWarnings("unchecked") void containerUpdated(RMContainer container, long updateTime); @SuppressWarnings("unchecked") void containerFinished(RMContainer container); @SuppressWarnings("unchecked") void computeQuota(ContainerId containerId, long computeTime);  }
@Test (timeout = 120000) public void TestQuotaService() throws Exception { int initialCredits = 50; int totalCost = 0; String user = "Project07__rizvi"; ApplicationId appId = ApplicationId.fromString("application_1450009406746_0001"); final List<ProjectQuota> hopProjectQuota = new ArrayList<ProjectQuota>(); hopProjectQuota.add(new ProjectQuota("Project07", initialCredits, 0)); LightWeightRequestHandler prepareHandler = new LightWeightRequestHandler( YARNOperationType.OTHER) { @Override public Object performTask() throws IOException { connector.beginTransaction(); connector.writeLock(); ProjectQuotaDataAccess<ProjectQuota> _pqDA = (ProjectQuotaDataAccess) RMStorageFactory. getDataAccess(ProjectQuotaDataAccess.class); _pqDA.addAll(hopProjectQuota); connector.commit(); return null; } }; prepareHandler.handle(); RMContext rmContext = Mockito.mock(RMContext.class); ConcurrentMap<ApplicationId, RMApp> rmApps = Mockito.mock(ConcurrentMap.class); RMApp rmApp = Mockito.mock(RMApp.class); PriceMultiplicatorService pms = Mockito.mock(PriceMultiplicatorService.class); Mockito.when(rmContext.getRMApps()).thenReturn(rmApps); Mockito.when(rmContext.getPriceMultiplicatorService()).thenReturn(pms); Mockito.when(rmApps.get(Mockito.any(ApplicationId.class))).thenReturn(rmApp); Mockito.when(rmApp.getUser()).thenReturn(user); Mockito.when(pms.getMultiplicator(Mockito.any(PriceMultiplicator.MultiplicatorType.class))).thenReturn(new Float(1)); QuotaService qs = new QuotaService(rmContext); conf.setInt(YarnConfiguration.QUOTA_SCHEDULING_PERIOD, 1200000); qs.init(conf); qs.serviceStart(); for (int i = 0; i < 10; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); RMContainer container = Mockito.mock(RMContainer.class); Mockito.when(container.getContainerId()).thenReturn(containerId); Mockito.when(container.getAllocatedResource()).thenReturn(TestUtils.createResource(1024, 1, ImmutableMap.<String, Integer>builder().build())); Mockito.when(container.getAllocatedNode()).thenReturn(null); Mockito.when(container.getAllocatedPriority()).thenReturn(Priority.UNDEFINED); Mockito.when(container.getCreationTime()).thenReturn(0L); qs.containerStarted(container); } } Thread.sleep(1000); for (int i = 0; i < 3; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); RMContainer container = Mockito.mock(RMContainer.class); Mockito.when(container.getContainerId()).thenReturn(containerId); Mockito.when(container.getDiagnosticsInfo()).thenReturn(null); Mockito.when(container.getContainerExitStatus()).thenReturn(0); Mockito.when(container.getContainerState()).thenReturn(ContainerState.COMPLETE); Mockito.when(container.getFinishTime()).thenReturn(1000L); qs.containerFinished(container); totalCost+=10; } } Thread.sleep(1000); for (int i = 3; i < 10; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); totalCost+=10; qs.computeQuota(containerId, 10000L); } } Thread.sleep(1000); for (int i = 3; i < 6; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); RMContainer container = Mockito.mock(RMContainer.class); Mockito.when(container.getContainerId()).thenReturn(containerId); Mockito.when(container.getDiagnosticsInfo()).thenReturn(null); Mockito.when(container.getContainerExitStatus()).thenReturn(0); Mockito.when(container.getContainerState()).thenReturn(ContainerState.COMPLETE); Mockito.when(container.getFinishTime()).thenReturn(19000L); qs.containerFinished(container); totalCost+=9; } } Thread.sleep(1000); for (int i = 6; i < 10; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); RMContainer container = Mockito.mock(RMContainer.class); Mockito.when(container.getContainerId()).thenReturn(containerId); Mockito.when(container.getAllocatedResource()).thenReturn(TestUtils.createResource(2048, 1, ImmutableMap.<String, Integer>builder().build())); Mockito.when(container.getAllocatedNode()).thenReturn(null); Mockito.when(container.getAllocatedPriority()).thenReturn(Priority.UNDEFINED); Mockito.when(container.getCreationTime()).thenReturn(0L); qs.containerUpdated(container, 20000L); totalCost+=10; } } Thread.sleep(1000); for (int i = 6; i < 10; i++) { ApplicationAttemptId appAttemptId = ApplicationAttemptId.newInstance(appId, i); for (int j = 0; j < i; j++) { ContainerId containerId = ContainerId.newContainerId(appAttemptId, j); RMContainer container = Mockito.mock(RMContainer.class); Mockito.when(container.getContainerId()).thenReturn(containerId); Mockito.when(container.getDiagnosticsInfo()).thenReturn(null); Mockito.when(container.getContainerExitStatus()).thenReturn(0); Mockito.when(container.getContainerState()).thenReturn(ContainerState.COMPLETE); Mockito.when(container.getFinishTime()).thenReturn(30000L); qs.containerFinished(container); totalCost+=20; } } Thread.sleep(1000); CheckProject(initialCredits - totalCost, totalCost); CheckProjectDailyCost(totalCost); }