/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.uber.athenax.backend.server.jobs; import com.uber.athenax.backend.api.InstanceStatus; import com.uber.athenax.backend.api.JobDefinition; import com.uber.athenax.backend.api.JobDefinitionDesiredstate; import com.uber.athenax.backend.api.JobDefinitionResource; import com.uber.athenax.backend.server.yarn.InstanceInfo; import com.uber.athenax.backend.server.yarn.InstanceMetadata; import org.junit.Test; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.UUID; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; public class JobWatcherUtilTest { @Test public void testComputeState() { HashMap<UUID, JobDefinition> jobMap = new HashMap<>(); HashMap<UUID, InstanceInfo> instanceMap = new HashMap<>(); mockState1(jobMap, instanceMap); JobWatcherUtil.StateView v = JobWatcherUtil.computeState(jobMap, instanceMap); HashMap<UUID, List<InstanceInfo>> jobInstances = new HashMap<>(); HashMap<UUID, UUID> instanceToJob = new HashMap<>(); for (Map.Entry<UUID, InstanceInfo> e : instanceMap.entrySet()) { UUID jobId = e.getValue().metadata().jobDefinition(); instanceToJob.put(e.getKey(), jobId); List<InstanceInfo> li = jobInstances.getOrDefault(jobId, new ArrayList<>()); li.add(e.getValue()); jobInstances.put(jobId, li); } assertEquals(jobMap, v.jobs()); assertEquals(instanceMap, v.instances()); assertEquals(instanceToJob, v.instanceToJob()); assertEquals(jobInstances, v.jobInstances()); } @Test public void testKillingSpuriousInstances() { UUID instId = UUID.randomUUID(); UUID jobId = UUID.randomUUID(); InstanceInfo instance = mockInstanceInfo("foo", 1, 4096, jobId, InstanceStatus.StateEnum.ACCEPTED); HealthCheckReport report = JobWatcherUtil.computeHealthCheckReport( Collections.emptyMap(), Collections.singletonMap(instId, instance)); assertEquals(1, report.spuriousInstances().size()); instance = mockInstanceInfo("foo", 1, 4096, jobId, InstanceStatus.StateEnum.KILLED); report = JobWatcherUtil.computeHealthCheckReport( Collections.emptyMap(), Collections.singletonMap(instId, instance)); assertTrue(report.spuriousInstances().isEmpty()); } @Test public void testHealthyInstance() { UUID instId = UUID.randomUUID(); UUID jobId = UUID.randomUUID(); InstanceInfo instance = mockInstanceInfo("foo", 1, 4096, jobId, InstanceStatus.StateEnum.RUNNING); JobDefinition job = mockJobWithSingleInstance("foo", 1, 4096); HealthCheckReport report = JobWatcherUtil.computeHealthCheckReport( Collections.singletonMap(jobId, job), Collections.singletonMap(instId, instance)); assertTrue(report.spuriousInstances().isEmpty()); assertTrue(report.instancesWithDifferentParameters().isEmpty()); assertTrue(report.instancesToStart().isEmpty()); } @Test public void testMismatchedParameter() { UUID instId = UUID.randomUUID(); UUID jobId = UUID.randomUUID(); InstanceInfo instance = mockInstanceInfo("foo", 1, 4096, jobId, InstanceStatus.StateEnum.RUNNING); JobDefinition job = mockJobWithSingleInstance("foo", 1, 2048); HealthCheckReport report = JobWatcherUtil.computeHealthCheckReport( Collections.singletonMap(jobId, job), Collections.singletonMap(instId, instance)); assertTrue(report.spuriousInstances().isEmpty()); assertEquals(1, report.instancesWithDifferentParameters().size()); assertEquals(1, report.instancesToStart().size()); } @Test public void testStartJob() { UUID jobId = UUID.randomUUID(); JobDefinition job = mockJobWithSingleInstance("foo", 1, 2048); HealthCheckReport report = JobWatcherUtil.computeHealthCheckReport( Collections.singletonMap(jobId, job), Collections.emptyMap()); assertTrue(report.spuriousInstances().isEmpty()); assertTrue(report.instancesWithDifferentParameters().isEmpty()); assertEquals(1, report.instancesToStart().size()); } private static JobDefinition mockJobWithSingleInstance(String clusterName, long vCore, long memory) { JobDefinitionDesiredstate desiredstate = mockDesiredState(clusterName, vCore, memory); JobDefinition job = mock(JobDefinition.class); doReturn(Collections.singletonList(desiredstate)).when(job).getDesiredState(); return job; } private static InstanceInfo mockInstanceInfo( String clusterName, long vCore, long memory, UUID jobId, InstanceStatus.StateEnum state) { InstanceInfo instance = mock(InstanceInfo.class); InstanceMetadata md = mock(InstanceMetadata.class); doReturn(md).when(instance).metadata(); doReturn(jobId).when(md).jobDefinition(); InstanceStatus status = mock(InstanceStatus.class); doReturn(state).when(status).getState(); doReturn(status).when(instance).status(); doReturn(vCore).when(status).getAllocatedVCores(); doReturn(memory).when(status).getAllocatedMB(); doReturn(clusterName).when(instance).clusterName(); return instance; } private static JobDefinitionDesiredstate mockDesiredState(String clusterName, long vCore, long memory) { return new JobDefinitionDesiredstate() .clusterId(clusterName) .resource(new JobDefinitionResource() .vCores(vCore) .memory(memory)); } private static void mockState1(Map<UUID, JobDefinition> jobs, Map<UUID, InstanceInfo> instances) { InstanceMetadata[] md = new InstanceMetadata[] { mock(InstanceMetadata.class), mock(InstanceMetadata.class), mock(InstanceMetadata.class) }; UUID[] instIds = new UUID[] { UUID.randomUUID(), UUID.randomUUID(), UUID.randomUUID() }; UUID[] appId = new UUID[3]; appId[0] = UUID.randomUUID(); appId[1] = appId[0]; appId[2] = UUID.randomUUID(); jobs.put(appId[0], mock(JobDefinition.class)); jobs.put(appId[2], mock(JobDefinition.class)); InstanceInfo[] instanceArray = new InstanceInfo[] { mock(InstanceInfo.class), mock(InstanceInfo.class), mock(InstanceInfo.class) }; for (int i = 0; i < instanceArray.length; ++i) { doReturn(md[i]).when(instanceArray[i]).metadata(); InstanceStatus status = mock(InstanceStatus.class); doReturn(InstanceStatus.StateEnum.RUNNING).when(status).getState(); doReturn(status).when(instanceArray[i]).status(); doReturn(instIds[i]).when(md[i]).uuid(); doReturn(appId[i]).when(md[i]).jobDefinition(); instances.put(instIds[i], instanceArray[i]); } } }