package com.hubspot.singularity.scheduler;
import static com.google.inject.name.Names.named;
import static com.hubspot.singularity.SingularityMainModule.HTTP_HOST_AND_PORT;
import static org.mockito.Mockito.*;
import java.util.Set;
import javax.servlet.http.HttpServletRequest;
import org.apache.curator.test.TestingServer;
import org.apache.mesos.Protos.MasterInfo;
import org.apache.mesos.Protos.Status;
import org.apache.mesos.SchedulerDriver;
import org.eclipse.jetty.util.component.LifeCycle;
import org.mockito.Matchers;
import org.slf4j.LoggerFactory;
import com.codahale.metrics.MetricRegistry;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Optional;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.google.common.net.HostAndPort;
import com.google.inject.Binder;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Module;
import com.google.inject.OutOfScopeException;
import com.google.inject.Provider;
import com.google.inject.Scopes;
import com.google.inject.Stage;
import com.google.inject.TypeLiteral;
import com.google.inject.util.Modules;
import com.hubspot.dropwizard.guicier.DropwizardModule;
import com.hubspot.dropwizard.guicier.GuiceBundle;
import com.hubspot.jackson.datatype.protobuf.ProtobufModule;
import com.hubspot.mesos.client.MesosClient;
import com.hubspot.singularity.SingularityAbort;
import com.hubspot.singularity.SingularityAuthModule;
import com.hubspot.singularity.SingularityMainModule;
import com.hubspot.singularity.SingularityTaskId;
import com.hubspot.singularity.SingularityTestAuthenticator;
import com.hubspot.singularity.auth.authenticator.SingularityAuthenticator;
import com.hubspot.singularity.config.MesosConfiguration;
import com.hubspot.singularity.config.SMTPConfiguration;
import com.hubspot.singularity.config.SentryConfiguration;
import com.hubspot.singularity.config.SingularityConfiguration;
import com.hubspot.singularity.config.ZooKeeperConfiguration;
import com.hubspot.singularity.data.SingularityDataModule;
import com.hubspot.singularity.data.history.SingularityHistoryModule;
import com.hubspot.singularity.data.transcoders.SingularityTranscoderModule;
import com.hubspot.singularity.data.zkmigrations.SingularityZkMigrationsModule;
import com.hubspot.singularity.event.SingularityEventModule;
import com.hubspot.singularity.hooks.LoadBalancerClient;
import com.hubspot.singularity.mesos.OfferCache;
import com.hubspot.singularity.mesos.SchedulerDriverSupplier;
import com.hubspot.singularity.mesos.SingularityDriver;
import com.hubspot.singularity.mesos.SingularityMesosExecutorInfoSupport;
import com.hubspot.singularity.mesos.SingularityMesosModule;
import com.hubspot.singularity.mesos.SingularityOfferCache;
import com.hubspot.singularity.resources.DeployResource;
import com.hubspot.singularity.resources.PriorityResource;
import com.hubspot.singularity.resources.RackResource;
import com.hubspot.singularity.resources.RequestResource;
import com.hubspot.singularity.resources.SlaveResource;
import com.hubspot.singularity.resources.TaskResource;
import com.hubspot.singularity.sentry.SingularityExceptionNotifier;
import com.hubspot.singularity.smtp.SingularityMailer;
import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import ch.qos.logback.classic.LoggerContext;
import io.dropwizard.db.DataSourceFactory;
import io.dropwizard.jackson.Jackson;
import io.dropwizard.setup.Environment;
import net.kencochrane.raven.Raven;
public class SingularityTestModule implements Module {
private final TestingServer ts;
private final DropwizardModule dropwizardModule;
private final ObjectMapper om = Jackson.newObjectMapper()
.setSerializationInclusion(Include.NON_NULL)
.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false)
.registerModule(new ProtobufModule());
private final Environment environment = new Environment("test-env", om, null, new MetricRegistry(), null);
private final boolean useDBTests;
public SingularityTestModule(boolean useDbTests) throws Exception {
this.useDBTests = useDbTests;
dropwizardModule = new DropwizardModule(environment);
LoggerContext context = (LoggerContext) LoggerFactory.getILoggerFactory();
Logger rootLogger = context.getLogger(org.slf4j.Logger.ROOT_LOGGER_NAME);
rootLogger.setLevel(Level.toLevel(System.getProperty("singularity.test.log.level", "WARN")));
Logger hsLogger = context.getLogger("com.hubspot");
hsLogger.setLevel(Level.toLevel(System.getProperty("singularity.test.log.level.for.com.hubspot", "WARN")));
this.ts = new TestingServer();
}
public Injector getInjector() throws Exception {
return Guice.createInjector(Stage.PRODUCTION, dropwizardModule, this);
}
public void start() throws Exception {
// Start all the managed instances in dropwizard.
Set<LifeCycle> managedObjects = ImmutableSet.copyOf(environment.lifecycle().getManagedObjects());
for (LifeCycle managed : managedObjects) {
managed.start();
}
}
public void stop() throws Exception {
ImmutableSet<LifeCycle> managedObjects = ImmutableSet.copyOf(environment.lifecycle().getManagedObjects());
for (LifeCycle managed : Lists.reverse(managedObjects.asList())) {
managed.stop();
}
}
@Override
public void configure(Binder mainBinder) {
mainBinder.install(new GuiceBundle.GuiceEnforcerModule());
TestingMesosClient tmc = new TestingMesosClient();
mainBinder.bind(MesosClient.class).toInstance(tmc);
mainBinder.bind(TestingMesosClient.class).toInstance(tmc);
mainBinder.bind(TestingServer.class).toInstance(ts);
final SingularityConfiguration configuration = getSingularityConfigurationForTestingServer(ts);
if (useDBTests) {
configuration.setDatabaseConfiguration(getDataSourceFactory());
}
mainBinder.bind(SingularityConfiguration.class).toInstance(configuration);
mainBinder.install(Modules.override(new SingularityMainModule(configuration))
.with(new Module() {
@Override
public void configure(Binder binder) {
binder.bind(SingularityExceptionNotifier.class).toInstance(mock(SingularityExceptionNotifier.class));
SingularityAbort abort = mock(SingularityAbort.class);
SingularityMailer mailer = mock(SingularityMailer.class);
binder.bind(SingularityMailer.class).toInstance(mailer);
binder.bind(SingularityAbort.class).toInstance(abort);
TestingLoadBalancerClient tlbc = new TestingLoadBalancerClient();
binder.bind(LoadBalancerClient.class).toInstance(tlbc);
binder.bind(TestingLoadBalancerClient.class).toInstance(tlbc);
binder.bind(OfferCache.class).to(SingularityOfferCache.class);
binder.bind(ObjectMapper.class).toInstance(om);
binder.bind(Environment.class).toInstance(environment);
binder.bind(HostAndPort.class).annotatedWith(named(HTTP_HOST_AND_PORT)).toInstance(HostAndPort.fromString("localhost:8080"));
binder.bind(new TypeLiteral<Optional<Raven>>() {}).toInstance(Optional.<Raven>absent());
binder.bind(new TypeLiteral<Optional<SentryConfiguration>>() {}).toInstance(Optional.<SentryConfiguration>absent());
binder.bind(HttpServletRequest.class).toProvider(new Provider<HttpServletRequest>() {
@Override
public HttpServletRequest get() {
throw new OutOfScopeException("testing");
}
});
}
}));
mainBinder.install(Modules.override(new SingularityMesosModule())
.with(new Module() {
@Override
public void configure(Binder binder) {
SingularityMesosExecutorInfoSupport logSupport = mock(SingularityMesosExecutorInfoSupport.class);
binder.bind(SingularityMesosExecutorInfoSupport.class).toInstance(logSupport);
SingularityDriver mock = mock(SingularityDriver.class);
when(mock.kill((SingularityTaskId) Matchers.any())).thenReturn(Status.DRIVER_RUNNING);
when(mock.getMaster()).thenReturn(Optional.<MasterInfo>absent());
when(mock.start()).thenReturn(Status.DRIVER_RUNNING);
when(mock.getLastOfferTimestamp()).thenReturn(Optional.<Long>absent());
binder.bind(SingularityDriver.class).toInstance(mock);
SchedulerDriver driver = mock(SchedulerDriver.class);
when(driver.killTask(null)).thenReturn(Status.DRIVER_RUNNING);
SchedulerDriverSupplier driverSupplier = new SchedulerDriverSupplier();
driverSupplier.setSchedulerDriver(driver);
binder.bind(SchedulerDriverSupplier.class).toInstance(driverSupplier);
}
}));
mainBinder.install(new SingularityDataModule());
mainBinder.install(new SingularitySchedulerModule());
mainBinder.install(new SingularityTranscoderModule());
mainBinder.install(new SingularityHistoryModule(configuration));
mainBinder.install(new SingularityZkMigrationsModule());
mainBinder.install(new SingularityEventModule(configuration));
mainBinder.install(Modules.override(new SingularityAuthModule(configuration))
.with(new Module() {
@Override
public void configure(Binder binder) {
binder.bind(SingularityAuthenticator.class).to(SingularityTestAuthenticator.class);
binder.bind(SingularityTestAuthenticator.class).in(Scopes.SINGLETON);
}
}));
mainBinder.bind(DeployResource.class);
mainBinder.bind(RequestResource.class);
mainBinder.bind(TaskResource.class);
mainBinder.bind(SlaveResource.class);
mainBinder.bind(RackResource.class);
mainBinder.bind(PriorityResource.class);
}
private DataSourceFactory getDataSourceFactory() {
DataSourceFactory dataSourceFactory = new DataSourceFactory();
dataSourceFactory.setDriverClass("org.h2.Driver");
dataSourceFactory.setUrl("jdbc:h2:mem:singularity;DB_CLOSE_DELAY=-1");
dataSourceFactory.setUser("user");
dataSourceFactory.setPassword("password");
return dataSourceFactory;
}
private static SingularityConfiguration getSingularityConfigurationForTestingServer(final TestingServer ts) {
SingularityConfiguration config = new SingularityConfiguration();
config.setLoadBalancerUri("test");
MesosConfiguration mc = new MesosConfiguration();
mc.setDefaultCpus(1);
mc.setDefaultMemory(128);
config.setMesosConfiguration(mc);
config.setSmtpConfiguration(new SMTPConfiguration());
ZooKeeperConfiguration zookeeperConfiguration = new ZooKeeperConfiguration();
zookeeperConfiguration.setQuorum(ts.getConnectString());
config.setZooKeeperConfiguration(zookeeperConfiguration);
config.setConsiderTaskHealthyAfterRunningForSeconds(0);
return config;
}
}