private void testEnableKerberos(final PrincipalKeyCredential PrincipalKeyCredential,
                                  String kdcType,
                                  String manageIdentities) throws Exception {

    StackId stackId = new StackId("HDP", "2.2");

    KerberosHelper kerberosHelper = injector.getInstance(KerberosHelper.class);
    boolean identitiesManaged = (manageIdentities == null) || !"false".equalsIgnoreCase(manageIdentities);

    final ServiceComponentHost schKerberosClient = createMock(ServiceComponentHost.class);
    expect(schKerberosClient.getServiceName()).andReturn(Service.Type.KERBEROS.name()).anyTimes();
    expect(schKerberosClient.getServiceComponentName()).andReturn(Role.KERBEROS_CLIENT.name()).anyTimes();
    expect(schKerberosClient.getHostName()).andReturn("host1").anyTimes();
    expect(schKerberosClient.getState()).andReturn(State.INSTALLED).anyTimes();

    final ServiceComponentHost sch1 = createMock(ServiceComponentHost.class);
    expect(sch1.getServiceName()).andReturn("SERVICE1").anyTimes();
    expect(sch1.getServiceComponentName()).andReturn("COMPONENT1").anyTimes();
    expect(sch1.getHostName()).andReturn("host1").anyTimes();
    expect(sch1.getState()).andReturn(State.INSTALLED).anyTimes();

    final ServiceComponentHost sch2 = createMock(ServiceComponentHost.class);
    expect(sch2.getServiceName()).andReturn("SERVICE2").anyTimes();
    expect(sch2.getServiceComponentName()).andReturn("COMPONENT2").anyTimes();
    expect(sch2.getHostName()).andReturn("host1").anyTimes();
    expect(sch2.getState()).andReturn(State.INSTALLED).anyTimes();

    final Host host = createMockHost("host1");

    final ServiceComponent serviceComponentKerberosClient = createNiceMock(ServiceComponent.class);
    expect(serviceComponentKerberosClient.getName()).andReturn(Role.KERBEROS_CLIENT.name()).anyTimes();
    expect(serviceComponentKerberosClient.getServiceComponentHosts()).andReturn(Collections.singletonMap("host1", schKerberosClient)).anyTimes();

    final Service serviceKerberos = createNiceMock(Service.class);
    expect(serviceKerberos.getDesiredStackId()).andReturn(stackId).anyTimes();
    expect(serviceKerberos.getName()).andReturn(Service.Type.KERBEROS.name()).anyTimes();
    expect(serviceKerberos.getServiceComponents())
        .andReturn(Collections.singletonMap(Role.KERBEROS_CLIENT.name(), serviceComponentKerberosClient))
        .anyTimes();

    final Service service1 = createNiceMock(Service.class);
    expect(service1.getDesiredStackId()).andReturn(stackId).anyTimes();
    expect(service1.getName()).andReturn("SERVICE1").anyTimes();
    expect(service1.getServiceComponents())
        .andReturn(Collections.emptyMap())
        .anyTimes();

    final Service service2 = createNiceMock(Service.class);
    expect(service2.getName()).andReturn("SERVICE2").anyTimes();
    expect(service2.getDesiredStackId()).andReturn(stackId).anyTimes();
    expect(service2.getServiceComponents())
        .andReturn(Collections.emptyMap())
        .anyTimes();

    final Map<String, String> kerberosEnvProperties = createMock(Map.class);
    expect(kerberosEnvProperties.get(KerberosHelper.KDC_TYPE)).andReturn(kdcType).anyTimes();
    expect(kerberosEnvProperties.get(KerberosHelper.MANAGE_IDENTITIES)).andReturn(manageIdentities).anyTimes();
    expect(kerberosEnvProperties.get(KerberosHelper.DEFAULT_REALM)).andReturn("FOOBAR.COM").anyTimes();
    expect(kerberosEnvProperties.get(KerberosHelper.CREATE_AMBARI_PRINCIPAL)).andReturn("false").anyTimes();
    expect(kerberosEnvProperties.get(KerberosHelper.PRECONFIGURE_SERVICES)).andReturn(PreconfigureServiceType.DEFAULT.name()).anyTimes();

    final Config kerberosEnvConfig = createMock(Config.class);
    expect(kerberosEnvConfig.getProperties()).andReturn(kerberosEnvProperties).anyTimes();

    final Map<String, String> krb5ConfProperties = createMock(Map.class);

    final Config krb5ConfConfig = createMock(Config.class);
    expect(krb5ConfConfig.getProperties()).andReturn(krb5ConfProperties).anyTimes();

    final Cluster cluster = createMockCluster("c1", Collections.singleton(host), SecurityType.KERBEROS, krb5ConfConfig, kerberosEnvConfig);
    expect(cluster.getDesiredStackVersion()).andReturn(new StackId("HDP-2.2")).anyTimes();
    expect(cluster.getServices())
        .andReturn(new HashMap<String, Service>() {
          {
            put(Service.Type.KERBEROS.name(), serviceKerberos);
            put("SERVICE1", service1);
            put("SERVICE2", service2);
          }
        })
        .anyTimes();
    expect(cluster.getServiceComponentHosts("host1"))
        .andReturn(new ArrayList<ServiceComponentHost>() {
          {
            add(schKerberosClient);
            add(sch1);
            add(sch2);
          }
        })
        .once();
    expect(cluster.getServiceComponentHosts("KERBEROS", "KERBEROS_CLIENT"))
        .andReturn(Collections.singletonList(schKerberosClient))
        .once();

    if (identitiesManaged) {
      final Clusters clusters = injector.getInstance(Clusters.class);
      expect(clusters.getHost("host1"))
          .andReturn(host)
          .once();
    }

    final AmbariManagementController ambariManagementController = injector.getInstance(AmbariManagementController.class);
    expect(ambariManagementController.findConfigurationTagsWithOverrides(cluster, null, null))
        .andReturn(Collections.emptyMap())
        .once();
    expect(ambariManagementController.getRoleCommandOrder(cluster))
        .andReturn(createMock(RoleCommandOrder.class))
        .once();

    final KerberosServiceDescriptor serviceDescriptor1 = createMock(KerberosServiceDescriptor.class);

    final KerberosServiceDescriptor serviceDescriptor2 = createMock(KerberosServiceDescriptor.class);

    final KerberosDescriptor kerberosDescriptor = createMock(KerberosDescriptor.class);
    expect(kerberosDescriptor.getService("KERBEROS")).andReturn(null).once();
    expect(kerberosDescriptor.getService("SERVICE1")).andReturn(serviceDescriptor1).once();
    expect(kerberosDescriptor.getService("SERVICE2")).andReturn(serviceDescriptor2).once();

    setupKerberosDescriptor(kerberosDescriptor);
    setupStageFactory();

    // This is a STRICT mock to help ensure that the end result is what we want.
    final RequestStageContainer requestStageContainer = createStrictMock(RequestStageContainer.class);
    // Create Preparation Stage
    expect(requestStageContainer.getLastStageId()).andReturn(-1L).anyTimes();
    expect(requestStageContainer.getId()).andReturn(1L).once();
    requestStageContainer.addStages(anyObject(List.class));
    expectLastCall().once();

    if (identitiesManaged) {
      // Create Principals Stage
      expect(requestStageContainer.getLastStageId()).andReturn(-1L).anyTimes();
      expect(requestStageContainer.getId()).andReturn(1L).once();
      requestStageContainer.addStages(anyObject(List.class));
      expectLastCall().once();
      // Create Keytabs Stage
      expect(requestStageContainer.getLastStageId()).andReturn(0L).anyTimes();
      expect(requestStageContainer.getId()).andReturn(1L).once();
      requestStageContainer.addStages(anyObject(List.class));
      expectLastCall().once();
      // Distribute Keytabs Stage
      expect(requestStageContainer.getLastStageId()).andReturn(1L).anyTimes();
      expect(requestStageContainer.getId()).andReturn(1L).once();
      requestStageContainer.addStages(anyObject(List.class));
      expectLastCall().once();
    }
    // Update Configs Stage
    expect(requestStageContainer.getLastStageId()).andReturn(2L).anyTimes();
    expect(requestStageContainer.getId()).andReturn(1L).once();

    requestStageContainer.addStages(anyObject(List.class));
    expectLastCall().once();
    // TODO: Add more of these when more stages are added.
    // Clean-up/Finalize Stage
    expect(requestStageContainer.getLastStageId()).andReturn(3L).anyTimes();
    expect(requestStageContainer.getId()).andReturn(1L).once();

    requestStageContainer.addStages(anyObject(List.class));
    expectLastCall().once();

    replayAll();

    // Needed by infrastructure
    injector.getInstance(AmbariMetaInfo.class).init();

    CredentialStoreService credentialStoreService = injector.getInstance(CredentialStoreService.class);
    credentialStoreService.setCredential(cluster.getClusterName(), KerberosHelper.KDC_ADMINISTRATOR_CREDENTIAL_ALIAS,
        PrincipalKeyCredential, CredentialStoreType.TEMPORARY);

    kerberosHelper.toggleKerberos(cluster, SecurityType.KERBEROS, requestStageContainer, null);

    verifyAll();
  }
Tests that when regenerating keytabs for an upgrade, that the retry allowed
boolean is set on the tasks created.

@throws Exception
ambari-server/src/test/java/org/apache/ambari/server/controller/KerberosHelperTest.java