private void prepareUpgrade() throws AmbariException {
    RepositoryVersionEntity repoVersionEntity = new RepositoryVersionEntity();
    repoVersionEntity.setDisplayName("Version to Upgrade To");
    repoVersionEntity.addRepoOsEntities(new ArrayList<>());
    repoVersionEntity.setStack(stackEntity220);
    repoVersionEntity.setVersion("2.2.0.1");
    repoVersionDAO.create(repoVersionEntity);

    helper.getOrCreateRepositoryVersion(stack220, stack220.getStackVersion());

    RequestEntity requestEntity = new RequestEntity();
    requestEntity.setRequestId(upgradeRequestId);
    requestEntity.setClusterId(cluster.getClusterId());
    requestDAO.create(requestEntity);

    // Create the stage and add it to the request
    StageEntity stageEntity = new StageEntity();
    stageEntity.setRequest(requestEntity);
    stageEntity.setClusterId(cluster.getClusterId());
    stageEntity.setRequestId(upgradeRequestId);
    stageEntity.setStageId(stageId);
    requestEntity.setStages(Collections.singletonList(stageEntity));
    stageDAO.create(stageEntity);
    requestDAO.merge(requestEntity);

    UpgradeEntity upgrade = new UpgradeEntity();
    upgrade.setId(1L);
    upgrade.setRequestEntity(requestEntity);
    upgrade.setClusterId(cluster.getClusterId());
    upgrade.setUpgradePackage("some-name");
    upgrade.setUpgradePackStackId(new StackId((String) null));
    upgrade.setUpgradeType(UpgradeType.ROLLING);
    upgrade.setDirection(Direction.UPGRADE);
    upgrade.setRepositoryVersion(repoVersionEntity);
    upgradeDAO.create(upgrade);

    cluster.setUpgradeEntity(upgrade);

    // Create the task and add it to the stage
    HostRoleCommandEntity hrc1 = new HostRoleCommandEntity();

    hrc1.setStage(stageEntity);
    hrc1.setStatus(HostRoleStatus.COMPLETED);
    hrc1.setRole(Role.ZOOKEEPER_SERVER);
    hrc1.setRoleCommand(RoleCommand.RESTART);

    stageEntity.setHostRoleCommands(new ArrayList<>());
    stageEntity.getHostRoleCommands().add(hrc1);
    hostRoleCommandDAO.create(hrc1);
    stageDAO.merge(stageEntity);
  }
Create a new repo version, plus the request and stage objects needed for a ROLLING stack upgrade.
@throws AmbariException
ambari-server/src/test/java/org/apache/ambari/server/state/services/RetryUpgradeActionServiceTest.java