src/test/java/org/apache/aurora/scheduler/thrift/ThriftIT.java

/**
 * 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.thrift;

import java.util.Optional;

import javax.inject.Singleton;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.Provides;
import com.google.inject.TypeLiteral;

import org.apache.aurora.common.application.ShutdownStage;
import org.apache.aurora.common.base.Command;
import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.apache.aurora.common.testing.easymock.EasyMockTest;
import org.apache.aurora.gen.AuroraAdmin;
import org.apache.aurora.gen.Container;
import org.apache.aurora.gen.Container._Fields;
import org.apache.aurora.gen.DockerContainer;
import org.apache.aurora.gen.DockerParameter;
import org.apache.aurora.gen.JobConfiguration;
import org.apache.aurora.gen.ScheduleStatus;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.ServerInfo;
import org.apache.aurora.gen.TaskConfig;
import org.apache.aurora.gen.TaskQuery;
import org.apache.aurora.scheduler.TierModule;
import org.apache.aurora.scheduler.app.AppModule;
import org.apache.aurora.scheduler.app.LifecycleModule;
import org.apache.aurora.scheduler.app.SchedulerMain;
import org.apache.aurora.scheduler.app.ServiceGroupMonitor;
import org.apache.aurora.scheduler.app.local.FakeNonVolatileStorage;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.config.CliOptions;
import org.apache.aurora.scheduler.config.types.TimeAmount;
import org.apache.aurora.scheduler.configuration.ConfigurationManager;
import org.apache.aurora.scheduler.configuration.ConfigurationManager.ConfigurationManagerSettings;
import org.apache.aurora.scheduler.configuration.executor.ExecutorSettings;
import org.apache.aurora.scheduler.cron.quartz.CronModule;
import org.apache.aurora.scheduler.maintenance.MaintenanceController;
import org.apache.aurora.scheduler.mesos.DriverFactory;
import org.apache.aurora.scheduler.mesos.DriverSettings;
import org.apache.aurora.scheduler.mesos.FrameworkInfoFactory;
import org.apache.aurora.scheduler.mesos.FrameworkInfoFactory.FrameworkInfoFactoryImpl;
import org.apache.aurora.scheduler.mesos.FrameworkInfoFactory.FrameworkInfoFactoryImpl.BaseFrameworkInfo;
import org.apache.aurora.scheduler.mesos.TestExecutorSettings;
import org.apache.aurora.scheduler.quota.QuotaModule;
import org.apache.aurora.scheduler.resources.ResourceTestUtil;
import org.apache.aurora.scheduler.resources.ResourceType;
import org.apache.aurora.scheduler.stats.StatsModule;
import org.apache.aurora.scheduler.storage.SnapshotStore;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.NonVolatileStorage;
import org.apache.aurora.scheduler.storage.backup.Recovery;
import org.apache.aurora.scheduler.storage.backup.StorageBackup;
import org.apache.aurora.scheduler.storage.entities.IResourceAggregate;
import org.apache.aurora.scheduler.storage.entities.IServerInfo;
import org.apache.aurora.scheduler.storage.mem.MemStorageModule;
import org.apache.aurora.scheduler.thrift.aop.AnnotatedAuroraAdmin;
import org.apache.aurora.scheduler.thrift.aop.AopModule;
import org.apache.mesos.v1.Protos.FrameworkInfo;
import org.apache.shiro.subject.Subject;
import org.junit.Test;

import static org.apache.aurora.gen.ResponseCode.OK;
import static org.apache.aurora.scheduler.app.SchedulerMain.Options.DriverKind.SCHEDULER_DRIVER;
import static org.junit.Assert.assertEquals;

public class ThriftIT extends EasyMockTest {

  private static final String USER = "someuser";
  private static final IResourceAggregate QUOTA = ResourceTestUtil.aggregate(1, 1, 1);
  private static final IServerInfo SERVER_INFO = IServerInfo.build(new ServerInfo());

  private AuroraAdmin.Iface thrift;

  private void createThrift(ConfigurationManagerSettings configurationManagerSettings) {
    ResourceType.initializeEmptyCliArgsForTest();
    CliOptions options = new CliOptions();
    Injector injector = Guice.createInjector(
        new ThriftModule(),
        new AopModule(options),
        new AbstractModule() {
          private <T> T bindMock(Class<T> clazz) {
            T mock = createMock(clazz);
            bind(clazz).toInstance(mock);
            return mock;
          }

          @Override
          protected void configure() {
            bind(CliOptions.class).toInstance(options);
            install(new LifecycleModule());
            install(new StatsModule(options.stats));
            install(new MemStorageModule());
            install(new QuotaModule());
            install(new CronModule(options.cron));
            install(new TierModule(TaskTestUtil.TIER_CONFIG));
            bind(ExecutorSettings.class).toInstance(TestExecutorSettings.THERMOS_EXECUTOR);

            install(new AppModule(configurationManagerSettings, SCHEDULER_DRIVER, options));
            install(new SchedulerMain.ProtocolModule(new SchedulerMain.Options()));

            bind(NonVolatileStorage.class).to(FakeNonVolatileStorage.class);

            ServiceGroupMonitor schedulers = createMock(ServiceGroupMonitor.class);
            bind(ServiceGroupMonitor.class).toInstance(schedulers);

            FrameworkInfo base = FrameworkInfo.newBuilder()
                    .setUser("framework user")
                    .setName("test framework")
                    .build();

            bindMock(DriverFactory.class);
            bind(DriverSettings.class).toInstance(new DriverSettings(
                "fakemaster",
                Optional.empty()));
            bind(FrameworkInfo.class)
                .annotatedWith(BaseFrameworkInfo.class)
                .toInstance(base);
            bind(FrameworkInfoFactory.class).to(FrameworkInfoFactoryImpl.class);
            bind(FrameworkInfoFactoryImpl.class).in(Singleton.class);
            bindMock(Recovery.class);
            bindMock(StorageBackup.class);
            bindMock(SnapshotStore.class);
            bind(IServerInfo.class).toInstance(SERVER_INFO);
            bind(new TypeLiteral<Amount<Long, Time>>() { })
                .annotatedWith(
                    MaintenanceController.MaintenanceControllerImpl.PollingInterval.class)
                .toInstance(new TimeAmount(1, Time.MINUTES));
          }

          @Provides
          Optional<Subject> provideSubject() {
            return Optional.of(createMock(Subject.class));
          }
        }
    );

    Command shutdownCommand =
        injector.getInstance(Key.get(Command.class, ShutdownStage.class));
    addTearDown(shutdownCommand::execute);

    thrift = injector.getInstance(AnnotatedAuroraAdmin.class);
    Storage storage = injector.getInstance(Storage.class);
    storage.prepare();
  }

  @Test
  public void testSetQuota() throws Exception {
    createThrift(TaskTestUtil.CONFIGURATION_MANAGER_SETTINGS);

    control.replay();

    assertEquals(
        OK,
        thrift.setQuota(USER, QUOTA.newBuilder()).getResponseCode());

    assertEquals(
        QUOTA.newBuilder(),
        thrift.getQuota(USER).getResult().getGetQuotaResult().getQuota());
  }

  @Test
  public void testSubmitNoExecutorDockerTask() throws Exception {
    ConfigurationManagerSettings configurationManagerSettings = new ConfigurationManagerSettings(
        ImmutableSet.of(_Fields.DOCKER),
        true,
        ImmutableList.of(),
        false,
        true,
        true,
        false,
        20,
        1800,
        ConfigurationManager.DEFAULT_ALLOWED_JOB_ENVIRONMENTS);

    createThrift(configurationManagerSettings);

    control.replay();

    TaskConfig task = TaskTestUtil.makeConfig(TaskTestUtil.JOB).newBuilder();
    task.unsetExecutorConfig();
    task.unsetSlaPolicy();
    task.setProduction(false)
        .setTier(TaskTestUtil.DEV_TIER_NAME)
        .setContainer(Container.docker(new DockerContainer()
            .setImage("image")
            .setParameters(
                ImmutableList.of(new DockerParameter("a", "b"), new DockerParameter("c", "d")))));
    JobConfiguration job = new JobConfiguration()
        .setKey(task.getJob())
        .setTaskConfig(task)
        .setInstanceCount(1);

    assertEquals(OK, thrift.createJob(job).getResponseCode());
    ScheduledTask scheduledTask = Iterables.getOnlyElement(
        thrift.getTasksStatus(new TaskQuery()).getResult().getScheduleStatusResult().getTasks());
    assertEquals(ScheduleStatus.PENDING, scheduledTask.getStatus());
    assertEquals(task, scheduledTask.getAssignedTask().getTask());
  }
}