com.spotify.helios.system.DeploymentGroupTest.java Source code

Java tutorial

Introduction

Here is the source code for com.spotify.helios.system.DeploymentGroupTest.java

Source

/*
 * Copyright (c) 2015 Spotify AB.
 *
 * 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 com.spotify.helios.system;

import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.spotify.helios.Polling;
import com.spotify.helios.agent.AgentMain;
import com.spotify.helios.client.HeliosClient;
import com.spotify.helios.common.Json;
import com.spotify.helios.common.descriptors.Deployment;
import com.spotify.helios.common.descriptors.DeploymentGroupStatus;
import com.spotify.helios.common.descriptors.Goal;
import com.spotify.helios.common.descriptors.HostStatus;
import com.spotify.helios.common.descriptors.Job;
import com.spotify.helios.common.descriptors.JobId;
import com.spotify.helios.common.descriptors.JobStatus;
import com.spotify.helios.common.descriptors.TaskStatus;
import com.spotify.helios.common.protocol.DeploymentGroupStatusResponse;
import com.spotify.helios.common.protocol.RemoveDeploymentGroupResponse;
import com.spotify.helios.common.protocol.RollingUpdateResponse;
import com.spotify.helios.master.MasterMain;

import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;

import static com.google.common.base.Strings.isNullOrEmpty;
import static com.google.common.collect.Iterables.getLast;
import static com.spotify.helios.common.descriptors.HostStatus.Status.UP;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.is;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThat;

public class DeploymentGroupTest extends SystemTestBase {

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private static final String TEST_GROUP = "my_group";
    private static final String TEST_LABEL = "foo=bar";
    private static final String TOKEN = "foo-token";

    private MasterMain master;

    @Before
    public void initialize() throws Exception {
        master = startDefaultMaster();

        // Wait for master to come up
        Polling.await(LONG_WAIT_SECONDS, SECONDS, new Callable<String>() {
            @Override
            public String call() throws Exception {
                final String output = cli("masters");
                return output.contains(masterName()) ? output : null;
            }
        });
    }

    @Test
    public void testListDeploymentGroups() throws Exception {
        cli("create-deployment-group", "group2", "foo=bar");
        cli("create-deployment-group", "group1", "foo=bar");
        final String output = cli("list-deployment-groups", "--json");
        final List<String> deploymentGroups = OBJECT_MAPPER.readValue(output, new TypeReference<List<String>>() {
        });
        assertEquals(Arrays.asList("group1", "group2"), deploymentGroups);
    }

    @Test
    public void testRollingUpdate() throws Exception {
        final List<String> hosts = ImmutableList.of("dc1-" + testHost() + "-a1.dc1.example.com",
                "dc1-" + testHost() + "-a2.dc1.example.com", "dc2-" + testHost() + "-a1.dc2.example.com",
                "dc2-" + testHost() + "-a3.dc2.example.com", "dc3-" + testHost() + "-a4.dc3.example.com");

        // start agents
        for (final String host : hosts) {
            startDefaultAgent(host, "--labels", TEST_LABEL);
        }

        // Wait for agents to come up
        final HeliosClient client = defaultClient();
        for (final String host : hosts) {
            awaitHostStatus(client, host, UP, LONG_WAIT_SECONDS, SECONDS);
        }

        // create a deployment group and job
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);

        // TODO: fix this!
        // Wait to make sure the host-update has run
        Thread.sleep(1000);

        // trigger a rolling update
        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);

        // ensure the job is running on all agents and the deployment group reaches DONE
        for (final String host : hosts) {
            awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
        }

        final Deployment deployment = defaultClient().hostStatus(hosts.get(0)).get().getJobs().get(jobId);
        assertEquals(TEST_GROUP, deployment.getDeploymentGroupName());
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        // create a second job
        final String secondJobVersion = testJobVersion + "2";
        final String secondJobNameAndVersion = testJobNameAndVersion + "2";
        final JobId secondJobId = createJob(testJobName, secondJobVersion, BUSYBOX, IDLE_COMMAND);

        // trigger a rolling update to replace the first job with the second job
        final String output = cli("rolling-update", secondJobNameAndVersion, TEST_GROUP);

        // Check that the hosts in the output are ordered
        final List<String> lines = Lists.newArrayList(Splitter.on("\n").split(output));
        for (int i = 0; i < hosts.size(); i++) {
            assertThat(lines.get(i + 2), containsString(hosts.get(i)));
        }

        // ensure the second job rolled out fine
        for (final String host : hosts) {
            awaitTaskState(secondJobId, host, TaskStatus.State.RUNNING);
        }
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
    }

    @Test
    public void testAgentAddedAfterRollingUpdateIsDeployed() throws Exception {
        startDefaultAgent(testHost(), "--labels", "foo=bar");

        cli("create-deployment-group", "--json", TEST_GROUP, "foo=bar");
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);

        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);

        awaitTaskState(jobId, testHost(), TaskStatus.State.RUNNING);

        // Rollout should be complete and on its second iteration at this point.
        // Start another agent and wait for it to have the job deployed to it.
        startDefaultAgent(testHost() + "2", "--labels", "foo=bar");

        awaitTaskState(jobId, testHost() + "2", TaskStatus.State.RUNNING);
    }

    @Test
    public void testRollingUpdateGroupNotFound() throws Exception {
        cli("create-deployment-group", "--json", TEST_GROUP, "foo=bar", "baz=qux");
        cli("create", "my_job:2", "my_image");
        assertEquals(RollingUpdateResponse.Status.DEPLOYMENT_GROUP_NOT_FOUND,
                OBJECT_MAPPER.readValue(cli("rolling-update", "--json", "--async", "my_job:2", "oops"),
                        RollingUpdateResponse.class).getStatus());
    }

    @Test
    public void testStatusNoRollingUpdate() throws Exception {
        cli("create-deployment-group", "--json", TEST_GROUP, "foo=bar", "baz=qux");
        assertEquals(DeploymentGroupStatusResponse.Status.IDLE,
                OBJECT_MAPPER.readValue(cli("deployment-group-status", "--json", TEST_GROUP),
                        DeploymentGroupStatusResponse.class).getStatus());
    }

    @Test
    public void testRollingUpdateMigrate() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Manually deploy a job on the host (i.e. a job not part of the deployment group)
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        deployJob(jobId, host);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);

        // Create a deployment-group and trigger a migration rolling-update
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--migrate", testJobNameAndVersion, TEST_GROUP);

        // Check that the deployment's deployment-group name eventually changes to TEST_GROUP
        // (should be null or empty before)
        final String jobDeploymentGroup = Polling.await(LONG_WAIT_SECONDS, SECONDS, new Callable<String>() {
            @Override
            public String call() throws Exception {
                final Deployment deployment = defaultClient().hostStatus(host).get().getJobs().get(jobId);
                if (deployment != null && !isNullOrEmpty(deployment.getDeploymentGroupName())) {
                    return deployment.getDeploymentGroupName();
                } else {
                    return null;
                }
            }
        });
        assertEquals(TEST_GROUP, jobDeploymentGroup);

        // rolling-update should succeed & job should be running
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
    }

    @Test
    public void testRollingUpdateMigrateWithToken() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Manually deploy a job with a token on the host (i.e. a job not part of the deployment group)
        final Job job = Job.newBuilder().setName(testJobName).setVersion(testJobVersion).setImage(BUSYBOX)
                .setCommand(IDLE_COMMAND).setToken(TOKEN).build();
        final JobId jobId = createJob(job);
        deployJob(jobId, host, TOKEN);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);

        // Create a deployment-group and trigger a migration rolling-update
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--migrate", "--token", TOKEN, testJobNameAndVersion, TEST_GROUP);

        // Check that the deployment's deployment-group name eventually changes to TEST_GROUP
        // (should be null or empty before)
        final String jobDeploymentGroup = Polling.await(LONG_WAIT_SECONDS, SECONDS, new Callable<String>() {
            @Override
            public String call() throws Exception {
                final Deployment deployment = defaultClient().hostStatus(host).get().getJobs().get(jobId);
                if (deployment != null && !isNullOrEmpty(deployment.getDeploymentGroupName())) {
                    return deployment.getDeploymentGroupName();
                } else {
                    return null;
                }
            }
        });
        assertEquals(TEST_GROUP, jobDeploymentGroup);

        // rolling-update should succeed & job should be running
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
    }

    @Test
    public void testRollingUpdateMigrateNothingToUndeploy() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Manually deploy a job on the host
        final String manualJobVersion = "foo-" + testJobVersion;
        final JobId manualJobId = createJob(testJobName, manualJobVersion, BUSYBOX, IDLE_COMMAND);
        deployJob(manualJobId, host);
        awaitTaskState(manualJobId, host, TaskStatus.State.RUNNING);

        // create a deployment group and trigger a migration rolling-update -- with a different
        // job that the one deployed manually! The manually deployed job should remain running on the
        // host.
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--migrate", testJobNameAndVersion, TEST_GROUP);

        // rolling-update should succeed & job should be running
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);

        final String jobDeploymentGroup = Polling.await(LONG_WAIT_SECONDS, SECONDS, new Callable<String>() {
            @Override
            public String call() throws Exception {
                final Deployment deployment = defaultClient().hostStatus(host).get().getJobs().get(jobId);
                if (deployment != null && !isNullOrEmpty(deployment.getDeploymentGroupName())) {
                    return deployment.getDeploymentGroupName();
                } else {
                    return null;
                }
            }
        });
        assertEquals(TEST_GROUP, jobDeploymentGroup);

        // Ensure that the manually deployed job is still there & running
        final Deployment manualDeployment = defaultClient().hostStatus(host).get().getJobs().get(manualJobId);
        assertNotNull(manualDeployment);
        assertEquals(Goal.START, manualDeployment.getGoal());
    }

    @Test
    public void testStopDeploymentGroup() throws Exception {
        cli("create-deployment-group", "--json", TEST_GROUP, "foo=bar", "baz=qux");
        cli("create", "my_job:2", "my_image");
        assertThat(cli("stop-deployment-group", TEST_GROUP), containsString("Deployment-group my_group stopped"));
        final DeploymentGroupStatusResponse status = Json.read(cli("deployment-group-status", "--json", TEST_GROUP),
                DeploymentGroupStatusResponse.class);
        assertEquals(DeploymentGroupStatusResponse.Status.FAILED, status.getStatus());
        assertEquals("Stopped by user", status.getError());
    }

    @Test
    public void testStopDeploymentGroupGroupNotFound() throws Exception {
        assertThat(cli("stop-deployment-group", TEST_GROUP), containsString("Deployment-group my_group not found"));
    }

    @Test
    public void testRollingUpdateCoordination() throws Exception {
        // stop the default master
        master.stopAsync().awaitTerminated();

        // start a bunch of masters and agents
        final Map<String, MasterMain> masters = startDefaultMasters(3);

        final Map<String, AgentMain> agents = Maps.newLinkedHashMap();
        for (int i = 0; i < 20; i++) {
            final String name = TEST_HOST + i;
            agents.put(name, startDefaultAgent(name, "--labels", TEST_LABEL));
        }

        // create a deployment group and start rolling out
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("rolling-update", "--async", "--par", String.valueOf(agents.size()), testJobNameAndVersion, TEST_GROUP);

        // wait until the task is running on the final agent
        awaitTaskState(jobId, getLast(agents.keySet()), TaskStatus.State.RUNNING);

        // ensure that all masters were involved
        final Set<String> deployingMasters = Sets.newHashSet();
        final Map<String, HostStatus> hostStatuses = defaultClient()
                .hostStatuses(Lists.newArrayList(agents.keySet())).get();
        for (final HostStatus status : hostStatuses.values()) {
            for (final Deployment deployment : status.getJobs().values()) {
                deployingMasters.add(deployment.getDeployerMaster());
            }
        }

        assertEquals(masters.size(), deployingMasters.size());
    }

    @Test
    public void testIdenticalRollouts() throws Exception {
        // This verifies that calling rolling-update on a failed deployment group will initiate a
        // new rollout and bring the group status out of the failed state. Previously calling
        // rolling-update on a failed deployment group would have no effect unless you changed the job
        // ID or a rollout option.

        // create the deployment group and job
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);

        // trigger a rolling update
        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        // stop the deployment group to put it in a failed state
        cli("stop-deployment-group", TEST_GROUP);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.FAILED);

        // trigger another rolling update with the same params as before and verify it reaches done
        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
    }

    @Test
    public void testStoppedJob() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Create the deployment group and two jobs
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        final JobId firstJobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        final String secondJobVersion = randomHexString();
        final String secondJobNameAndVersion = testJobName + ":" + secondJobVersion;
        final JobId secondJobId = createJob(testJobName, secondJobVersion, BUSYBOX, IDLE_COMMAND);

        // Trigger a rolling update of the first job
        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);
        awaitTaskState(firstJobId, host, TaskStatus.State.RUNNING);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        // Stop the job
        cli("stop", testJobNameAndVersion, host);
        awaitTaskState(firstJobId, host, TaskStatus.State.STOPPED);

        // Trigger a rolling update, replacing the first job with the second.
        // Verify the first job is undeployed and the second job is running.
        cli("rolling-update", "--async", secondJobNameAndVersion, TEST_GROUP);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(secondJobId, host, TaskStatus.State.RUNNING);
        final JobStatus status = client.jobStatus(firstJobId).get();
        assertThat(status.getDeployments().isEmpty(), is(true));

        // Stop the job
        cli("stop", secondJobNameAndVersion, host);
        awaitTaskState(secondJobId, host, TaskStatus.State.STOPPED);

        // Trigger a rolling update of the same job, and verify the job gets started. This takes
        // a different code path than when replacing a different job, which we tested above.
        cli("rolling-update", "--async", secondJobNameAndVersion, TEST_GROUP);
        awaitTaskState(secondJobId, host, TaskStatus.State.RUNNING);
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
    }

    @Test
    public void testRollingUpdateWithOverlap() throws Exception {
        final List<String> hosts = ImmutableList.of("dc1-" + testHost() + "-a1.dc1.example.com",
                "dc1-" + testHost() + "-a2.dc1.example.com", "dc2-" + testHost() + "-a1.dc2.example.com",
                "dc2-" + testHost() + "-a3.dc2.example.com", "dc3-" + testHost() + "-a4.dc3.example.com");
        // start agents
        for (final String host : hosts) {
            startDefaultAgent(host, "--labels", TEST_LABEL);
        }

        // create a deployment group
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);

        // create and roll out a first job
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("rolling-update", "--async", "--overlap", testJobNameAndVersion, TEST_GROUP);

        for (final String host : hosts) {
            awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
        }
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        // create and roll out a second job
        final String secondJobVersion = testJobVersion + "2";
        final String secondJobNameAndVersion = testJobNameAndVersion + "2";
        final JobId secondJobId = createJob(testJobName, secondJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("rolling-update", "--async", "--overlap", secondJobNameAndVersion, TEST_GROUP);

        for (final String host : hosts) {
            awaitTaskState(secondJobId, host, TaskStatus.State.RUNNING);
        }
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
    }

    @Ignore
    @Test
    public void testRollingUpdatePerformance() throws Exception {
        final List<String> hosts = ImmutableList.of("dc1-" + testHost() + "-a1.dc1.example.com",
                "dc1-" + testHost() + "-a2.dc1.example.com", "dc2-" + testHost() + "-a1.dc2.example.com",
                "dc2-" + testHost() + "-a3.dc2.example.com", "dc3-" + testHost() + "-a4.dc3.example.com");

        // start agents
        for (final String host : hosts) {
            startDefaultAgent(host, "--labels", TEST_LABEL);
        }

        // Wait for agents to come up
        final HeliosClient client = defaultClient();
        for (final String host : hosts) {
            awaitHostStatus(client, host, UP, LONG_WAIT_SECONDS, SECONDS);
        }

        for (int i = 0; i < 50; ++i) {
            cli("create-deployment-group", "--json", TEST_GROUP + "-" + i, "tol=ahdsksajd");
            cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP + "-" + i);
        }

        // create a deployment group and job
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);

        // TODO: fix this!
        // Wait for the host-updater
        Thread.sleep(2000);

        // trigger a rolling update
        cli("rolling-update", "--async", testJobNameAndVersion, TEST_GROUP);

        final long t0 = System.currentTimeMillis();

        // ensure the job is running on all agents and the deployment group reaches DONE
        for (final String host : hosts) {
            awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
        }

        final Deployment deployment = defaultClient().hostStatus(hosts.get(0)).get().getJobs().get(jobId);
        assertEquals(TEST_GROUP, deployment.getDeploymentGroupName());
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        System.out.printf("1 active / 0 inactive: Time to roll out: %.2f s\n",
                (System.currentTimeMillis() - t0) / 1000.0);
    }

    @Test
    public void testRollingUpdateWithOverlapAndParallelism() throws Exception {
        // create and start agents
        final List<String> hosts = ImmutableList.of("dc1-" + testHost() + "-a1.dc1.example.com",
                "dc1-" + testHost() + "-a2.dc1.example.com", "dc2-" + testHost() + "-a1.dc2.example.com",
                "dc2-" + testHost() + "-a3.dc2.example.com", "dc3-" + testHost() + "-a4.dc3.example.com");
        for (final String host : hosts) {
            startDefaultAgent(host, "--labels", TEST_LABEL);
        }

        // create a deployment group
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);

        // create and roll out a first job
        final JobId jobId = createJob(testJobName, testJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("rolling-update", "--async", "-p", "2", "--overlap", testJobNameAndVersion, TEST_GROUP);

        for (final String host : hosts) {
            awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
        }
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);

        // create and roll out a second job
        final String secondJobVersion = testJobVersion + "2";
        final String secondJobNameAndVersion = testJobNameAndVersion + "2";
        final JobId secondJobId = createJob(testJobName, secondJobVersion, BUSYBOX, IDLE_COMMAND);
        cli("rolling-update", "--async", "-p", "2", "--overlap", secondJobNameAndVersion, TEST_GROUP);

        for (final String host : hosts) {
            awaitTaskState(secondJobId, host, TaskStatus.State.RUNNING);
        }
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
    }

    @Test
    public void testRollingUpdateWithToken() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Manually deploy a job with a token on the host (i.e. a job not part of the deployment group)
        final Job job = Job.newBuilder().setName(testJobName).setVersion(testJobVersion).setImage(BUSYBOX)
                .setCommand(IDLE_COMMAND).setToken(TOKEN).build();
        final JobId jobId = createJob(job);

        // Create a deployment-group and trigger a migration rolling-update
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--token", TOKEN, testJobNameAndVersion, TEST_GROUP);

        // rolling-update should succeed & job should be running
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);

        // Check that we cannot manually undeploy the job with a token
        final String output = cli("undeploy", jobId.toString(), host);
        assertThat(output, containsString("FORBIDDEN"));
        awaitDeploymentGroupStatus(defaultClient(), TEST_GROUP, DeploymentGroupStatus.State.DONE);
        awaitTaskState(jobId, host, TaskStatus.State.RUNNING);
    }

    @Test
    public void testRemoveDeploymentGroupInactive() throws Exception {
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--migrate", testJobNameAndVersion, TEST_GROUP);

        final RemoveDeploymentGroupResponse response = defaultClient().removeDeploymentGroup(TEST_GROUP).get();

        assertEquals(RemoveDeploymentGroupResponse.Status.REMOVED, response.getStatus());
    }

    @Test
    public void testRemoveDeploymentGroupActive() throws Exception {
        final String host = testHost();
        startDefaultAgent(host, "--labels", TEST_LABEL);

        // Wait for agent to come up
        final HeliosClient client = defaultClient();
        awaitHostStatus(client, testHost(), UP, LONG_WAIT_SECONDS, SECONDS);

        // Use an invalid image to make sure the DG doesn't succeed (this ensure the DG will be active
        // for 5+ minutes, which is plenty of time).
        final JobId jobId = createJob(testJobName, testJobVersion, "invalid_image", IDLE_COMMAND);
        cli("create-deployment-group", "--json", TEST_GROUP, TEST_LABEL);
        cli("rolling-update", "--async", "--migrate", testJobNameAndVersion, TEST_GROUP);

        final RemoveDeploymentGroupResponse response = defaultClient().removeDeploymentGroup(TEST_GROUP).get();

        assertEquals(RemoveDeploymentGroupResponse.Status.REMOVED, response.getStatus());
    }
}