Merge pull request #217 from spotify/rculbertson/tmp-job-prober

Custom probers can now be provided to each job

helios-testing/src/main/java/com/spotify/helios/testing/DefaultDeployer.java

@@ -41,25 +41,23 @@
 import static java.util.Arrays.asList;
 import static org.junit.Assert.fail;
 
-public class DefaultDeployer implements TemporaryJob.Deployer {
+public class DefaultDeployer implements Deployer {
 
   private static final Logger log = LoggerFactory.getLogger(DefaultDeployer.class);
 
   private final HeliosClient client;
   private final List<TemporaryJob> jobs;
-  private final Prober prober;
 
   private boolean readyToDeploy;
 
-  public DefaultDeployer(HeliosClient client, List<TemporaryJob> jobs, Prober prober) {
+  public DefaultDeployer(HeliosClient client, List<TemporaryJob> jobs) {
     this.client = client;
     this.jobs = jobs;
-    this.prober = prober;
   }
 
   @Override
   public TemporaryJob deploy(final Job job, final String hostFilter,
-                             final Set<String> waitPorts) {
+                             final Set<String> waitPorts, final Prober prober) {
     if (isNullOrEmpty(hostFilter)) {
       fail("a host filter pattern must be passed to hostFilter(), " +
            "or one must be specified in HELIOS_HOST_FILTER");
@@ -83,12 +81,12 @@ public TemporaryJob deploy(final Job job, final String hostFilter,
     }
 
     final String chosenHost = filteredHosts.get(new Random().nextInt(filteredHosts.size()));
-    return deploy(job, asList(chosenHost), waitPorts);
+    return deploy(job, asList(chosenHost), waitPorts, prober);
   }
 
   @Override
   public TemporaryJob deploy(final Job job, final List<String> hosts,
-                             final Set<String> waitPorts) {
+                             final Set<String> waitPorts, final Prober prober) {
     if (!readyToDeploy) {
       fail("deploy() must be called in a @Before or in the test method, or perhaps you forgot"
            + " to put @Rule before TemporaryJobs");

helios-testing/src/main/java/com/spotify/helios/testing/Deployer.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2014 Spotify AB.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 com.spotify.helios.testing;
+
+import com.spotify.helios.common.descriptors.Job;
+
+import java.util.List;
+import java.util.Set;
+
+public interface Deployer {
+
+  TemporaryJob deploy(Job job, List<String> hosts, Set<String> waitPorts, Prober prober);
+
+  TemporaryJob deploy(Job job, String hostFilter, Set<String> waitPorts, Prober prober);
+
+  void readyToDeploy();
+
+}

helios-testing/src/main/java/com/spotify/helios/testing/TemporaryJob.java

@@ -305,12 +305,4 @@ private String endpointFromHost(String host) {
     return ip == null ? host : ip;
   }
 
-  public static interface Deployer {
-
-    TemporaryJob deploy(Job job, List<String> hosts, Set<String> waitPorts);
-
-    TemporaryJob deploy(Job job, String hostFilter, Set<String> waitPorts);
-
-    void readyToDeploy();
-  }
 }

helios-testing/src/main/java/com/spotify/helios/testing/TemporaryJobBuilder.java

@@ -50,6 +50,7 @@
 import static com.fasterxml.jackson.databind.node.JsonNodeType.STRING;
 import static com.google.common.base.Charsets.UTF_8;
 import static com.google.common.base.Optional.fromNullable;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static java.lang.Integer.toHexString;
 import static java.lang.System.getenv;
@@ -64,16 +65,21 @@ public class TemporaryJobBuilder {
   private final List<String> hosts = Lists.newArrayList();
   private final Job.Builder builder = Job.newBuilder();
   private final Set<String> waitPorts = Sets.newHashSet();
-  private final TemporaryJob.Deployer deployer;
+  private final Deployer deployer;
   private final String jobNamePrefix;
 
-  private String hostFilter = null;
-
+  private String hostFilter;
+  private Prober prober;
   private TemporaryJob job;
 
-  public TemporaryJobBuilder(final TemporaryJob.Deployer deployer, final String jobNamePrefix) {
+  public TemporaryJobBuilder(final Deployer deployer, final String jobNamePrefix,
+                             final Prober defaultProber) {
+    checkNotNull(deployer, "deployer");
+    checkNotNull(jobNamePrefix, "jobNamePrefix");
+    checkNotNull(defaultProber, "defaultProber");
     this.deployer = deployer;
     this.jobNamePrefix = jobNamePrefix;
+    this.prober = defaultProber;
     this.builder.setRegistrationDomain(jobNamePrefix);
   }
 
@@ -177,6 +183,16 @@ public TemporaryJobBuilder expires(final Date expires) {
     return this;
   }
 
+  /**
+   * This will override the default prober provided by {@link TemporaryJobs} to the constructor.
+   * @param prober the prober to use for this job
+   * @return the TemporaryJobBuilder
+   */
+  public TemporaryJobBuilder prober(final Prober prober) {
+    this.prober = prober;
+    return this;
+  }
+
   /**
    * Deploys the job to the specified hosts. If no hosts are specified, a host will be chosen at
    * random from the current Helios cluster. If the HELIOS_HOST_FILTER environment variable is set,
@@ -220,9 +236,9 @@ public TemporaryJob deploy(final List<String> hosts) {
           hostFilter = getenv("HELIOS_HOST_FILTER");
         }
 
-        job = deployer.deploy(builder.build(), hostFilter, waitPorts);
+        job = deployer.deploy(builder.build(), hostFilter, waitPorts, prober);
       } else {
-        job = deployer.deploy(builder.build(), this.hosts, waitPorts);
+        job = deployer.deploy(builder.build(), this.hosts, waitPorts, prober);
       }
     }
 

helios-testing/src/main/java/com/spotify/helios/testing/TemporaryJobs.java

@@ -73,7 +73,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 
 public class TemporaryJobs implements TestRule {
-  private static final Logger log = LoggerFactory.getLogger(TemporaryJob.class);
+  private static final Logger log = LoggerFactory.getLogger(TemporaryJobs.class);
 
   static final String HELIOS_TESTING_PROFILE = "helios.testing.profile";
   private static final String HELIOS_TESTING_PROFILES = "helios.testing.profiles.";
@@ -90,7 +90,7 @@ public class TemporaryJobs implements TestRule {
   private final JobPrefixFile jobPrefixFile;
   private final Config config;
   private final List<TemporaryJob> jobs = Lists.newCopyOnWriteArrayList();
-  private TemporaryJob.Deployer deployer;
+  private final Deployer deployer;
 
   private final ExecutorService executor = MoreExecutors.getExitingExecutorService(
       (ThreadPoolExecutor) Executors.newFixedThreadPool(
@@ -104,8 +104,7 @@ public class TemporaryJobs implements TestRule {
     this.client = checkNotNull(builder.client, "client");
     this.prober = checkNotNull(builder.prober, "prober");
     this.defaultHostFilter = checkNotNull(builder.hostFilter, "hostFilter");
-    this.deployer = Optional.fromNullable(builder.deployer)
-        .or(new DefaultDeployer(client, jobs, prober));
+    this.deployer = Optional.fromNullable(builder.deployer).or(new DefaultDeployer(client, jobs));
     final Path prefixDirectory = Paths.get(Optional.fromNullable(builder.prefixDirectory)
         .or(DEFAULT_PREFIX_DIRECTORY));
 
@@ -166,7 +165,8 @@ public void after() {
   }
 
   public TemporaryJobBuilder job() {
-    final TemporaryJobBuilder builder = new TemporaryJobBuilder(deployer, jobPrefixFile.prefix());
+    final TemporaryJobBuilder builder = new TemporaryJobBuilder(deployer, jobPrefixFile.prefix(),
+                                                                prober);
 
     if (config.hasPath("env")) {
       final Config env = config.getConfig("env");
@@ -492,7 +492,7 @@ private Builder(final String profile, final Config preConfig) {
     private final Config config;
     private String user = DEFAULT_USER;
     private Prober prober = DEFAULT_PROBER;
-    private TemporaryJob.Deployer deployer;
+    private Deployer deployer;
     private String hostFilter = DEFAULT_HOST_FILTER;
     private HeliosClient client;
     private String prefixDirectory;
@@ -536,7 +536,7 @@ public Builder prober(final Prober prober) {
       return this;
     }
 
-    public Builder deployer(final TemporaryJob.Deployer deployer) {
+    public Builder deployer(final Deployer deployer) {
       this.deployer = deployer;
       return this;
     }

helios-testing/src/test/java/com/spotify/helios/testing/ConfigTest.java

@@ -45,7 +45,7 @@ public void setup() {
     when(client.jobs()).thenReturn(future);
   }
 
-  public static class ProfileTest implements TemporaryJob.Deployer {
+  public static class ProfileTest implements Deployer {
 
     // Local is the default profile, so don't specify it explicitly to test default loading
     @Rule
@@ -69,15 +69,15 @@ public void test() throws Exception {
     }
 
     @Override
-    public TemporaryJob deploy(Job job, List<String> hosts, Set<String> waitPorts) {
+    public TemporaryJob deploy(Job job, List<String> hosts, Set<String> waitPorts, Prober prober) {
       // This is called when the first job is deployed
       assertThat(hosts, equalTo((List<String>) newArrayList("test-host")));
       parameters.validate(job, temporaryJobs.prefix());
       return null;
     }
 
     @Override
-    public TemporaryJob deploy(Job job, String hostFilter, Set<String> waitPorts) {
+    public TemporaryJob deploy(Job job, String hostFilter, Set<String> waitPorts, Prober prober) {
       // This is called when the second job is deployed
       assertThat(hostFilter, equalTo(parameters.hostFilter));
       parameters.validate(job, temporaryJobs.prefix());

helios-testing/src/test/java/com/spotify/helios/testing/ProberTest.java

@@ -0,0 +1,79 @@
+package com.spotify.helios.testing;
+
+import com.spotify.helios.client.HeliosClient;
+import com.spotify.helios.system.SystemTestBase;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import static com.spotify.helios.common.descriptors.HostStatus.Status.UP;
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+import static org.junit.experimental.results.PrintableResult.testResult;
+import static org.junit.experimental.results.ResultMatchers.isSuccessful;
+
+public class ProberTest extends SystemTestBase {
+
+  private static HeliosClient client;
+  private static String testHost;
+
+  public static class OverrideDefaultProberTest {
+
+    private MockProber defaultProber = new MockProber();
+    private MockProber overrideProber = new MockProber();
+
+    @Rule
+    public final TemporaryJobs temporaryJobs = TemporaryJobs.builder()
+        .client(client)
+        .prober(defaultProber)
+        .build();
+
+    @Before
+    public void setup() {
+      temporaryJobs.job()
+          .command(IDLE_COMMAND)
+          .port("default", 4711)
+          .deploy(testHost);
+
+      temporaryJobs.job()
+          .command(IDLE_COMMAND)
+          .port("override", 4712)
+          .prober(overrideProber)
+          .deploy(testHost);
+    }
+
+    @Test
+    public void test() {
+      // Verify that the first job used the prober passed to the TemporaryJobs rule.
+      assertThat(defaultProber.probed(), is(true));
+      // Verify that the second job used the prober that was passed to its builder.
+      assertThat(overrideProber.probed(), is(true));
+    }
+  }
+
+  @Test
+  public void testOverrideDefaultProber() throws Exception {
+    startDefaultMaster();
+    client = defaultClient();
+    testHost = testHost();
+    startDefaultAgent(testHost);
+    awaitHostStatus(client, testHost, UP, LONG_WAIT_MINUTES, MINUTES);
+    assertThat(testResult(OverrideDefaultProberTest.class), isSuccessful());
+  }
+
+  private static class MockProber implements Prober {
+    private boolean probed;
+
+    @Override
+    public boolean probe(String host, int port) {
+      return probed = true;
+    }
+
+    public boolean probed() {
+      return probed;
+    }
+  }
+
+}