@Test
public void testAutomaitcLogAggregationDelegationToken()
    throws Exception {

  Configuration conf = getConf();
  SecurityUtil.setAuthenticationMethod(
      UserGroupInformation.AuthenticationMethod.KERBEROS, conf);
  conf.set(YarnConfiguration.RM_PRINCIPAL, YARN_RM);
  String remoteRootLogPath = "/tmp/app-logs";
  MiniDFSCluster hdfsCluster = null;
  try {
    // Step 1: Start a MiniDFSCluster for Log Aggregation Path
    HdfsConfiguration hdfsConfig = new HdfsConfiguration();
    hdfsCluster = new MiniDFSCluster.Builder(hdfsConfig)
        .numDataNodes(1).build();
    Path remoteRootLogDir = new Path(remoteRootLogPath);
    FileSystem fs = hdfsCluster.getFileSystem();
    fs.mkdirs(remoteRootLogDir);
    conf.set(YarnConfiguration.NM_REMOTE_APP_LOG_DIR,
        fs.getFileStatus(remoteRootLogDir).getPath().toString());
    // Step 2: Prepare a Mock FileSystem which returns Delegation Token
    // when YarnClientImpl invokes
    DelegationTokenIdentifier hdfsDT = new DelegationTokenIdentifier(new Text(
        "test"), new Text(YARN_RM), null);
    final Token<DelegationTokenIdentifier> dToken =
        new Token<>(hdfsDT.getBytes(), new byte[0], hdfsDT.getKind(),
        new Text());
    FileSystem mockFs = mock(FileSystem.class);
    doAnswer(new Answer<Token<?>[]>() {
      @Override
      public Token<?>[] answer(InvocationOnMock invocation) {
        Object[] args = invocation.getArguments();
        ((Credentials) args[1]).addToken(hdfsDT.getKind(), dToken);
        return new Token[]{dToken};
      }
    }).when(mockFs).addDelegationTokens(any(), any());
    FileSystemTestHelper.addFileSystemForTesting(fs.getUri(),
        hdfsConfig, mockFs);
    // Step 3: Prepare a Mock YarnClientImpl
    YarnClientImpl client = spy(new YarnClientImpl() {
      @Override
      protected void serviceStart() {
        rmClient = mock(ApplicationClientProtocol.class);
      }
      @Override
      protected void serviceStop() {
      }
      @Override
      public ApplicationReport getApplicationReport(ApplicationId appId) {
        ApplicationReport report = mock(ApplicationReport.class);
        when(report.getYarnApplicationState())
            .thenReturn(YarnApplicationState.RUNNING);
        return report;
      }
      @Override
      public boolean isSecurityEnabled() {
        return true;
      }
    });
    client.init(conf);
    client.start();
    // Step 4: Prepare a ApplicationSubmissionContext and submit the app
    ApplicationSubmissionContext context =
        mock(ApplicationSubmissionContext.class);
    ApplicationId applicationId = ApplicationId.newInstance(0, 1);
    when(context.getApplicationId()).thenReturn(applicationId);
    DataOutputBuffer dob = new DataOutputBuffer();
    Credentials credentials = new Credentials();
    credentials.writeTokenStorageToStream(dob);
    ByteBuffer tokens = ByteBuffer.wrap(dob.getData(), 0, dob.getLength());
    ContainerLaunchContext clc = ContainerLaunchContext.newInstance(
        null, null, null, null, tokens, null);
    when(context.getAMContainerSpec()).thenReturn(clc);
    client.submitApplication(context);
    // Step 5: Verify automatic addition of HDFS DT for log aggregation path
    credentials = new Credentials();
    DataInputByteBuffer dibb = new DataInputByteBuffer();
    tokens = clc.getTokens();
    if (tokens != null) {
      dibb.reset(tokens);
      credentials.readTokenStorageStream(dibb);
      tokens.rewind();
    }
    Collection<Token<? extends TokenIdentifier>> dTokens =
         credentials.getAllTokens();
    Assert.assertEquals("Failed to place token for Log Aggregation Path",
        1, dTokens.size());
    Assert.assertEquals("Wrong Token for Log Aggregation",
        hdfsDT.getKind(), dTokens.iterator().next().getKind());
  } finally {
    if (hdfsCluster != null) {
      hdfsCluster.shutdown();
    }
  }
}