package com.hubspot.singularity.scheduler;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import org.apache.mesos.Protos.FrameworkID;
import org.apache.mesos.Protos.Offer;
import org.apache.mesos.Protos.OfferID;
import org.apache.mesos.Protos.Resource;
import org.apache.mesos.Protos.SlaveID;
import org.apache.mesos.Protos.Value.Range;
import org.apache.mesos.Protos.Value.Ranges;
import org.apache.mesos.Protos.Value.Scalar;
import org.apache.mesos.Protos.Value.Type;
import org.junit.Assert;
import org.junit.Test;
import com.google.common.base.Optional;
import com.google.common.collect.Lists;
import com.hubspot.mesos.MesosUtils;
import com.hubspot.singularity.ExtendedTaskState;
import com.hubspot.singularity.SingularityTaskHistoryUpdate;
import com.hubspot.singularity.SingularityTaskId;
public class MesosUtilsTest {
private void assertFound(int numPorts, Resource resource) {
int portsFound = 0;
for (Range r : resource.getRanges().getRangeList()) {
portsFound += (r.getEnd() - r.getBegin()) + 1;
}
Assert.assertEquals(numPorts, portsFound);
}
private void test(int numPorts, String... ranges) {
Resource resource = MesosUtils.getPortsResource(numPorts, buildOffer(ranges));
assertFound(numPorts, resource);
}
@Test
public void testTaskOrdering() {
final SingularityTaskId taskId = new SingularityTaskId("r", "d", System.currentTimeMillis(), 1, "h", "r");
final Optional<String> msg = Optional.absent();
SingularityTaskHistoryUpdate update1 = new SingularityTaskHistoryUpdate(taskId, 1L, ExtendedTaskState.TASK_LAUNCHED, msg);
SingularityTaskHistoryUpdate update2 = new SingularityTaskHistoryUpdate(taskId, 2L, ExtendedTaskState.TASK_RUNNING, msg);
SingularityTaskHistoryUpdate update3 = new SingularityTaskHistoryUpdate(taskId, 2L, ExtendedTaskState.TASK_FAILED, msg);
List<SingularityTaskHistoryUpdate> list = Arrays.asList(update2, update1, update3);
Collections.sort(list);
Assert.assertTrue(list.get(0).getTaskState() == ExtendedTaskState.TASK_LAUNCHED);
Assert.assertTrue(list.get(1).getTaskState() == ExtendedTaskState.TASK_RUNNING);
Assert.assertTrue(list.get(2).getTaskState() == ExtendedTaskState.TASK_FAILED);
}
@Test
public void testSubtractResources() {
Assert.assertEquals(createResources(3, 60, "23:23", "100:175", "771:1000"),
MesosUtils.subtractResources(createResources(5, 100, "23:23", "100:1000"), createResources(2, 40, "176:770")));
List<Resource> subtracted = createResources(100, 1000, "1:100", "101:1000");
subtracted = MesosUtils.subtractResources(subtracted, createResources(5, 100, "23:74", "101:120", "125:130", "750:756"));
Assert.assertEquals(createResources(95, 900, "1:22", "75:100", "121:124", "131:749", "757:1000"), subtracted);
subtracted = MesosUtils.subtractResources(subtracted, createResources(20, 20, "75:90", "121:121", "757:1000"));
Assert.assertEquals(createResources(75, 880, "1:22", "91:100", "122:124", "131:749"), subtracted);
}
private List<Resource> createResources(int cpus, int memory, String... ranges) {
List<Resource> resources = Lists.newArrayList();
if (cpus > 0) {
resources.add(Resource.newBuilder().setType(Type.SCALAR).setName(MesosUtils.CPUS).setScalar(Scalar.newBuilder().setValue(cpus).build()).build());
}
if (memory > 0) {
resources.add(Resource.newBuilder().setType(Type.SCALAR).setName(MesosUtils.MEMORY).setScalar(Scalar.newBuilder().setValue(memory).build()).build());
}
if (ranges.length > 0) {
resources.add(buildRanges(ranges));
}
return resources;
}
@Test
public void testRangeSelection() {
test(4, "23:24", "26:26", "28:28", "29:29", "31:32");
test(2, "22:23");
test(3, "22:22", "23:23", "24:24", "25:25");
test(10, "100:10000");
test(23, "90:100", "9100:9100", "185:1000");
}
private Resource buildRanges(String... ranges) {
Resource.Builder resources = Resource.newBuilder()
.setType(Type.RANGES)
.setName(MesosUtils.PORTS);
Ranges.Builder rangesBuilder = Ranges.newBuilder();
for (String range : ranges) {
String[] split = range.split("\\:");
rangesBuilder.addRange(
Range.newBuilder()
.setBegin(Long.parseLong(split[0]))
.setEnd(Long.parseLong(split[1])));
}
resources.setRanges(rangesBuilder);
return resources.build();
}
private Offer buildOffer(String... ranges) {
Offer.Builder offer = Offer.newBuilder()
.setId(OfferID.newBuilder().setValue("offerid").build())
.setFrameworkId(FrameworkID.newBuilder().setValue("frameworkid").build())
.setHostname("hostname")
.setSlaveId(SlaveID.newBuilder().setValue("slaveid").build());
offer.addResources(buildRanges(ranges));
return offer.build();
}
}