@Test
    public void successGetChangeHostModeStatus() throws Throwable {

      TaskState returnedState = new TaskState();
      returnedState.stage = TaskState.TaskStage.FINISHED;

      ChangeHostModeTaskService.State returnedDocument = new ChangeHostModeTaskService.State();
      returnedDocument.hostServiceLink = "host-id";
      returnedDocument.hostMode = HostMode.ENTERING_MAINTENANCE;
      returnedDocument.documentSelfLink = "task-id";
      returnedDocument.taskState = returnedState;

      setupMock(
          host,
          true,
          "operationId",
          ChangeHostModeTaskService.State.class,
          (state) -> {
            assertThat(state.hostServiceLink, is(HostServiceFactory.SELF_LINK + "/" + "host-id"));
            return true;
          },
          returnedDocument);

      TaskState state = target.getChangeHostModeStatus("operationId");
      assertThat(state.stage, is(TaskState.TaskStage.FINISHED));
    }
    @Test
    public void successChangeHostMode() throws Throwable {
      EnterMaintenanceModeRequest request = createEnterMaintenanceModeRequest();

      ChangeHostModeTaskService.State returnedDocument = new ChangeHostModeTaskService.State();
      returnedDocument.hostServiceLink = "host-id";
      returnedDocument.hostMode = HostMode.ENTERING_MAINTENANCE;
      returnedDocument.documentSelfLink = "task-id";

      setupMock(
          host,
          true,
          ChangeHostModeTaskFactoryService.SELF_LINK,
          ChangeHostModeTaskService.State.class,
          (state) -> {
            assertThat(state.hostServiceLink, is(HostServiceFactory.SELF_LINK + "/" + "host-id"));
            return true;
          },
          returnedDocument);

      String taskLink = target.changeHostMode(request.getHostId(), HostMode.ENTERING_MAINTENANCE);
      assertThat(taskLink, is("task-id"));
    }