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.util.Optional;
  
  import com.google.common.collect.ImmutableList;
  
  import org.apache.aurora.common.testing.easymock.EasyMockTest;
  import org.apache.aurora.scheduler.events.PubsubEvent;
  import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
  import org.apache.mesos.v1.Protos;
  import org.apache.mesos.v1.scheduler.Mesos;
  import org.apache.mesos.v1.scheduler.Protos.Call;
  import org.apache.mesos.v1.scheduler.Scheduler;
  import org.easymock.Capture;
  import org.junit.Before;
  import org.junit.Test;
  
  import static org.easymock.EasyMock.capture;
  import static org.easymock.EasyMock.expect;
  import static org.easymock.EasyMock.expectLastCall;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertTrue;
  
  public class VersionedSchedulerDriverServiceTest extends EasyMockTest {
  
    private static final String FRAMEWORK_ID = "framework_id";
    private static final Protos.OfferID OFFER_ID =
        Protos.OfferID.newBuilder().setValue("offer-id").build();
    private static final Protos.Filters FILTER =
        Protos.Filters.newBuilder().setRefuseSeconds(10).build();
    private static final Protos.FrameworkInfo BASE_INFO = Protos.FrameworkInfo.newBuilder()
            .setUser("framework user")
            .setName("test framework")
            .build();
    private static final DriverSettings SETTINGS = new DriverSettings(
        "fakemaster",
        Optional.empty());
  
    private StorageTestUtil storage;
    private Scheduler scheduler;
    private VersionedSchedulerDriverService driverService;
    private VersionedDriverFactory driverFactory;
    private FrameworkInfoFactory infoFactory;
    private Mesos mesos;
  
    @Before
    public void setUp() {
      scheduler = createMock(Scheduler.class);
      storage = new StorageTestUtil(this);
      driverFactory = createMock(VersionedDriverFactory.class);
      mesos = createMock(Mesos.class);
      infoFactory = createMock(FrameworkInfoFactory.class);
      driverService = new VersionedSchedulerDriverService(
          storage.storage,
          SETTINGS,
          scheduler,
          driverFactory,
          infoFactory);
    }
  
    @Test
    public void testNoopStop() {
      control.replay();
  
      driverService.stopAsync().awaitTerminated();
    }
  
    @Test
    public void testStop() {
      expectStart();
      control.replay();
  
      driverService.startAsync().awaitRunning();
      driverService.stopAsync().awaitTerminated();
    }
  
    @Test(expected = IllegalStateException.class)
    public void testExceptionBeforeStart() {
      control.replay();
      driverService.killTask("task-id");
    }
  
    @Test
    public void testBlockingBeforeRegistered() throws InterruptedException {
      expectStart();
      control.replay();
      driverService.startAsync().awaitRunning();
  
      Thread killRunner = new Thread(() -> {
        driverService.killTask("task-id");
      });
  
      killRunner.start();
  
      // A hack to ensure the thread actually executes the method
      Thread.sleep(1000L);
      assertEquals(Thread.State.WAITING, killRunner.getState());
  
      killRunner.interrupt();
    }
  
    @Test
    public void testKill() {
      expectStart();
      expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));
  
      Capture<Call> killCapture = createCapture();
      mesos.send(capture(killCapture));
      expectLastCall().once();
  
      control.replay();
      driverService.startAsync().awaitRunning();
      driverService.registered(new PubsubEvent.DriverRegistered());
  
      driverService.killTask("task-id");
  
      assertTrue(killCapture.hasCaptured());
      Call kill = killCapture.getValue();
      assertEquals(kill.getFrameworkId().getValue(), FRAMEWORK_ID);
      assertEquals(kill.getType(), Call.Type.KILL);
      assertEquals(kill.getKill().getTaskId().getValue(), "task-id");
    }
  
    @Test
    public void testDecline() {
      expectStart();
      expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));
  
      Capture<Call> declineCapture = createCapture();
      mesos.send(capture(declineCapture));
      expectLastCall().once();
  
      control.replay();
      driverService.startAsync().awaitRunning();
      driverService.registered(new PubsubEvent.DriverRegistered());
  
      driverService.declineOffer(OFFER_ID, FILTER);
  
      assertTrue(declineCapture.hasCaptured());
      Call decline = declineCapture.getValue();
      assertEquals(decline.getFrameworkId().getValue(), FRAMEWORK_ID);
      assertEquals(decline.getType(), Call.Type.DECLINE);
      assertEquals(decline.getDecline().getOfferIdsList(), ImmutableList.of(OFFER_ID));
      assertEquals(decline.getDecline().getFilters(), FILTER);
    }
  
    @Test
    public void testAccept() {
      expectStart();
      expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));
  
      Capture<Call> acceptCapture = createCapture();
      mesos.send(capture(acceptCapture));
      expectLastCall().once();
  
      control.replay();
      driverService.startAsync().awaitRunning();
      driverService.registered(new PubsubEvent.DriverRegistered());
  
      driverService.acceptOffers(OFFER_ID, ImmutableList.of(), FILTER);
  
      assertTrue(acceptCapture.hasCaptured());
      Call accept = acceptCapture.getValue();
      assertEquals(accept.getFrameworkId().getValue(), FRAMEWORK_ID);
      assertEquals(accept.getType(), Call.Type.ACCEPT);
      assertEquals(accept.getAccept().getFilters(), FILTER);
      assertEquals(accept.getAccept().getOfferIdsList(), ImmutableList.of(OFFER_ID));
      assertEquals(accept.getAccept().getOperationsList(), ImmutableList.of());
    }
  
    @Test
    public void testAcceptInverseOffer() {
      expectStart();
  
      expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));
  
      Capture<Call> acceptCapture = createCapture();
      mesos.send(capture(acceptCapture));
      expectLastCall().once();
  
      control.replay();
      driverService.startAsync().awaitRunning();
      driverService.registered(new PubsubEvent.DriverRegistered());
  
      driverService.acceptInverseOffer(OFFER_ID, FILTER);
  
      assertTrue(acceptCapture.hasCaptured());
      Call accept = acceptCapture.getValue();
      assertEquals(accept.getFrameworkId().getValue(), FRAMEWORK_ID);
      assertEquals(accept.getType(), Call.Type.ACCEPT_INVERSE_OFFERS);
      assertEquals(accept.getAcceptInverseOffers().getFilters(), FILTER);
      assertEquals(
          accept.getAcceptInverseOffers().getInverseOfferIdsList(),
          ImmutableList.of(OFFER_ID)
      );
    }
  
    private void expectStart() {
      storage.expectOperations();
      expect(storage.schedulerStore.fetchFrameworkId()).andReturn(Optional.of(FRAMEWORK_ID));
      expect(infoFactory.getFrameworkInfo()).andReturn(BASE_INFO);
  
      Protos.FrameworkInfo.Builder builder = BASE_INFO.toBuilder();
      builder.setId(Protos.FrameworkID.newBuilder().setValue(FRAMEWORK_ID));
  
      expect(driverFactory.create(
          scheduler,
          builder.build(),
          SETTINGS.getMasterUri(),
          SETTINGS.getCredentials()
      )).andReturn(mesos);
    }
  }
  