private UpgradeContext getMockUpgradeContext(Cluster cluster, Direction direction,
      UpgradeType type, RepositoryVersionEntity repositoryVersion, final RepositoryType repositoryType,
      Set<String> services, MasterHostResolver resolver, boolean replay) {
    UpgradeContext context = EasyMock.createNiceMock(UpgradeContext.class);
    expect(context.getCluster()).andReturn(cluster).anyTimes();
    expect(context.getType()).andReturn(type).anyTimes();
    expect(context.getDirection()).andReturn(direction).anyTimes();
    expect(context.getRepositoryVersion()).andReturn(repositoryVersion).anyTimes();
    expect(context.getSupportedServices()).andReturn(services).anyTimes();
    expect(context.getOrchestrationType()).andReturn(repositoryType).anyTimes();
    expect(context.getAmbariMetaInfo()).andReturn(ambariMetaInfo).anyTimes();
    expect(context.getHostRoleCommandFactory()).andStubReturn(
        injector.getInstance(HostRoleCommandFactory.class));
    expect(context.getRoleGraphFactory()).andStubReturn(
        injector.getInstance(RoleGraphFactory.class));

    // only set this if supplied
    if (null != resolver) {
      expect(context.getResolver()).andReturn(resolver).anyTimes();
    }

    final Map<String, RepositoryVersionEntity> targetRepositoryVersions = new HashMap<>();
    for( String serviceName : services ){
      targetRepositoryVersions.put(serviceName, repositoryVersion);
    }

    final Capture<String> repoVersionServiceName = EasyMock.newCapture();
    expect(
        context.getTargetRepositoryVersion(EasyMock.capture(repoVersionServiceName))).andAnswer(
            new IAnswer<RepositoryVersionEntity>() {
              @Override
              public RepositoryVersionEntity answer() {
                return targetRepositoryVersions.get(repoVersionServiceName.getValue());
              }
            }).anyTimes();

    final Capture<String> serviceNameSupported = EasyMock.newCapture();
    expect(context.isServiceSupported(EasyMock.capture(serviceNameSupported))).andAnswer(
        new IAnswer<Boolean>() {
          @Override
          public Boolean answer() {
            return targetRepositoryVersions.containsKey(serviceNameSupported.getValue());
          }
        }).anyTimes();


    final Map<String, String> serviceNames = new HashMap<>();
    final Capture<String> serviceDisplayNameArg1 = EasyMock.newCapture();
    final Capture<String> serviceDisplayNameArg2 = EasyMock.newCapture();

    context.setServiceDisplay(EasyMock.capture(serviceDisplayNameArg1), EasyMock.capture(serviceDisplayNameArg2));
    expectLastCall().andAnswer(
        new IAnswer<Object>() {
          @Override
          public Object answer() {
            serviceNames.put(serviceDisplayNameArg1.getValue(), serviceDisplayNameArg2.getValue());
            return null;
          }
        }).anyTimes();


    final Map<String, String> componentNames = new HashMap<>();
    final Capture<String> componentDisplayNameArg1 = EasyMock.newCapture();
    final Capture<String> componentDisplayNameArg2 = EasyMock.newCapture();
    final Capture<String> componentDisplayNameArg3 = EasyMock.newCapture();

    context.setComponentDisplay(EasyMock.capture(componentDisplayNameArg1),
        EasyMock.capture(componentDisplayNameArg2), EasyMock.capture(componentDisplayNameArg3));

    expectLastCall().andAnswer(new IAnswer<Object>() {
      @Override
      public Object answer() {
        componentNames.put(
            componentDisplayNameArg1.getValue() + ":" + componentDisplayNameArg2.getValue(),
            componentDisplayNameArg3.getValue());
        return null;
      }
    }).anyTimes();

    final Capture<String> getServiceDisplayArgument1 = EasyMock.newCapture();
    expect(
        context.getServiceDisplay(EasyMock.capture(getServiceDisplayArgument1))).andAnswer(
            new IAnswer<String>() {
              @Override
              public String answer() {
                return serviceNames.get(getServiceDisplayArgument1.getValue());
              }
            }).anyTimes();

    final Capture<String> getComponentDisplayArgument1 = EasyMock.newCapture();
    final Capture<String> getComponentDisplayArgument2 = EasyMock.newCapture();
    expect(context.getComponentDisplay(EasyMock.capture(getComponentDisplayArgument1),
        EasyMock.capture(getComponentDisplayArgument2))).andAnswer(new IAnswer<String>() {
          @Override
          public String answer() {
            return componentNames.get(getComponentDisplayArgument1.getValue() + ":"
                + getComponentDisplayArgument2.getValue());
          }
        }).anyTimes();

    final Capture<UpgradeScope> isScopedCapture = EasyMock.newCapture();
    expect(context.isScoped(EasyMock.capture(isScopedCapture))).andStubAnswer(
        new IAnswer<Boolean>() {
          @Override
          public Boolean answer() throws Throwable {
            UpgradeScope scope = isScopedCapture.getValue();
            if (scope == UpgradeScope.ANY) {
              return true;
            }

            if (scope == UpgradeScope.PARTIAL) {
              return repositoryType != RepositoryType.STANDARD;
            }

            return repositoryType == RepositoryType.STANDARD;
          }
        });

    if (replay) {
      replay(context);
    }

    return context;
  }
Builds a mock upgrade context using only the supplied method arguments.

@param cluster
         the cluster for the upgrade
@param direction
         the upgrade direction
@param type
         the type of upgrade
@param repositoryVersion
         the repository version to use for the upgrade or downgrade.
@param repositoryType
         the type of repository for the upgrade (patch, standard, etc).
@param services
         the services participating in the upgrade. This should typically
         be all services, unless the {@link RepositoryType#PATCH} is used
@param resolver
         the master hsot resolver to use when determining how to
         orchestrate hosts.
@param replay
         {@code true} to replay the mock object before returning it,
         {@code false} otherwise.
@return a mock {@link UpgradeContext} which has already been replayed and
        is ready to use.
ambari-server/src/test/java/org/apache/ambari/server/stack/upgrade/orchestrate/UpgradeHelperTest.java