content: /**
   * Licensed 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 org.apache.aurora.scheduler.mesos;
  
  import java.nio.charset.StandardCharsets;
  
  import com.google.common.collect.ImmutableList;
  import com.google.common.net.InetAddresses;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.mesos.Protos;
  import org.apache.mesos.SchedulerDriver;
  import org.apache.mesos.v1.Protos.ExecutorID;
  import org.apache.mesos.v1.Protos.FrameworkID;
  import org.apache.mesos.v1.Protos.OfferID;
  import org.apache.mesos.v1.Protos.TaskID;
  import org.apache.mesos.v1.Protos.TaskState;
  import org.apache.mesos.v1.Protos.TaskStatus;
  import org.apache.mesos.v1.Protos.TaskStatus.Reason;
  import org.apache.mesos.v1.Protos.TaskStatus.Source;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.apache.aurora.scheduler.mesos.ProtosConversion.convert;
  import static org.easymock.EasyMock.expectLastCall;
  
  public class MesosSchedulerImplTest extends EasyMockTest {
  
    private static final String FRAMEWORK_ID = "framework-id";
    private static final FrameworkID FRAMEWORK =
        FrameworkID.newBuilder().setValue(FRAMEWORK_ID).build();
    private static final String MASTER_ID = "master-id";
    private static final Protos.MasterInfo MASTER = Protos.MasterInfo.newBuilder()
        .setId(MASTER_ID)
        .setIp(InetAddresses.coerceToInteger(InetAddresses.forString("1.2.3.4"))) //NOPMD
        .setPort(5050).build();
    private static final String SLAVE_HOST = "slave-hostname";
    private static final Protos.SlaveID SLAVE_ID =
        Protos.SlaveID.newBuilder().setValue("slave-id").build();
  
    private static final OfferID OFFER_ID = OfferID.newBuilder().setValue("offer-id").build();
    private static final Protos.Offer OFFER = Protos.Offer.newBuilder()
            .setFrameworkId(convert(FRAMEWORK))
            .setSlaveId(SLAVE_ID)
            .setHostname(SLAVE_HOST)
            .setId(convert(OFFER_ID))
            .build();
  
    private static final ExecutorID EXECUTOR_ID =
        ExecutorID.newBuilder().setValue("executor-id").build();
  
    private static final TaskStatus STATUS_NO_REASON = TaskStatus.newBuilder()
        .setState(TaskState.TASK_RUNNING)
        .setSource(Source.SOURCE_AGENT)
        .setMessage("message")
        .setTimestamp(1D)
        .setTaskId(TaskID.newBuilder().setValue("task-id").build())
        .build();
  
    private static final TaskStatus STATUS = STATUS_NO_REASON
        .toBuilder()
        // Only testing data plumbing, this field with TASK_RUNNING would not normally happen,
        .setReason(Reason.REASON_COMMAND_EXECUTOR_FAILED)
        .build();
  
    private MesosCallbackHandler handler;
    private SchedulerDriver driver;
  
    private MesosSchedulerImpl scheduler;
  
    @Before
    public void setUp() {
      handler = createMock(MesosCallbackHandler.class);
      driver = createMock(SchedulerDriver.class);
      scheduler = new MesosSchedulerImpl(handler);
    }
  
    @Test
    public void testSlaveLost() {
      handler.handleLostAgent(convert(SLAVE_ID));
      expectLastCall().once();
  
      control.replay();
  
      scheduler.slaveLost(driver, SLAVE_ID);
    }
  
    @Test
    public void testRegistered() {
      handler.handleRegistration(FRAMEWORK, convert(MASTER));
      expectLastCall().once();
  
      control.replay();
  
      scheduler.registered(driver, convert(FRAMEWORK), MASTER);
    }
  
    @Test
    public void testDisconnected() {
      handler.handleDisconnection();
      expectLastCall().once();
  
      control.replay();
  
      scheduler.disconnected(driver);
    }
  
    @Test
    public void testReRegistered() {
      handler.handleReregistration(convert(MASTER));
      expectLastCall().once();
  
      control.replay();
  
      scheduler.reregistered(driver, MASTER);
    }
  
    @Test
    public void testResourceOffers() {
      handler.handleRegistration(FRAMEWORK, convert(MASTER));
      expectLastCall().once();
      handler.handleOffers(ImmutableList.of(convert(OFFER)));
      expectLastCall().once();
  
      control.replay();
  
      scheduler.registered(driver, convert(FRAMEWORK), MASTER);
      scheduler.resourceOffers(driver, ImmutableList.of(OFFER));
    }
  
    @Test(expected = IllegalStateException.class)
    public void testBadOrdering() {
      control.replay();
  
      // Should fail since the scheduler is not yet registered.
      scheduler.resourceOffers(driver, ImmutableList.of());
    }
  
    @Test
    public void testOfferRescinded() {
      handler.handleRescind(OFFER_ID);
      expectLastCall().once();
  
      control.replay();
  
      scheduler.offerRescinded(driver, convert(OFFER_ID));
    }
  
    @Test
    public void testStatusUpdate() {
      handler.handleUpdate(STATUS);
      expectLastCall().once();
  
      control.replay();
  
      scheduler.statusUpdate(driver, convert(STATUS));
    }
  
    @Test
    public void testError() {
      handler.handleError("Oh No!");
      expectLastCall().once();
  
      control.replay();
  
      scheduler.error(driver, "Oh No!");
    }
  
    @Test
    public void testFrameworkMessage() {
      handler.handleMessage(EXECUTOR_ID, convert(SLAVE_ID));
      expectLastCall().once();
  
      control.replay();
  
      scheduler.frameworkMessage(
          driver,
          convert(EXECUTOR_ID),
          SLAVE_ID,
          "message".getBytes(StandardCharsets.UTF_8));
    }
  
    @Test
    public void testExecutorLost() {
      handler.handleLostExecutor(EXECUTOR_ID, convert(SLAVE_ID), 1);
  
      control.replay();
  
      scheduler.executorLost(driver, convert(EXECUTOR_ID), SLAVE_ID, 1);
    }
  }
  