@VisibleForTesting
    void setTokensFor(ContainerLaunchContext containerLaunchContext, boolean fetchToken)
            throws Exception {
        Credentials credentials = new Credentials();

        LOG.info("Loading delegation tokens available locally to add to the AM container");
        // for user
        UserGroupInformation currUsr = UserGroupInformation.getCurrentUser();

        Collection<Token<? extends TokenIdentifier>> usrTok =
                currUsr.getCredentials().getAllTokens();
        for (Token<? extends TokenIdentifier> token : usrTok) {
            LOG.info("Adding user token " + token.getService() + " with " + token);
            credentials.addToken(token.getService(), token);
        }

        if (fetchToken) {
            LOG.info("Fetching delegation tokens to add to the AM container.");
            DelegationTokenManager delegationTokenManager =
                    new DefaultDelegationTokenManager(flinkConfiguration, null, null, null);
            DelegationTokenContainer container = new DelegationTokenContainer();
            delegationTokenManager.obtainDelegationTokens(container);

            // This is here for backward compatibility to make log aggregation work
            for (Map.Entry<String, byte[]> e : container.getTokens().entrySet()) {
                if (flinkConfiguration.get(APP_MASTER_TOKEN_SERVICES).contains(e.getKey())) {
                    credentials.addAll(HadoopDelegationTokenConverter.deserialize(e.getValue()));
                }
            }
        }

        ByteBuffer tokens = ByteBuffer.wrap(HadoopDelegationTokenConverter.serialize(credentials));
        containerLaunchContext.setTokens(tokens);

        LOG.info("Delegation tokens added to the AM container.");
    }
@Test
    public void testSetTokensForYarnAppMaster() {
        final Configuration flinkConfig = new Configuration();
        flinkConfig.set(
                APP_MASTER_TOKEN_SERVICES,
                Arrays.asList(TestYarnAMDelegationTokenProvider.SERVICE_NAME));
        YarnClusterDescriptor yarnClusterDescriptor = createYarnClusterDescriptor(flinkConfig);
        ContainerLaunchContext amContainer = Records.newRecord(ContainerLaunchContext.class);
        try {
            yarnClusterDescriptor.setTokensFor(amContainer, true);
            Credentials credentials = new Credentials();
            try (DataInputStream dis =
                    new DataInputStream(
                            new ByteArrayInputStream(amContainer.getTokens().array()))) {
                credentials.readTokenStorageStream(dis);
            }
            assertThat(credentials.getAllTokens())
                    .hasSize(1)
                    .contains(TestYarnAMDelegationTokenProvider.TEST_YARN_AM_TOKEN);
        } catch (Exception e) {
            fail("Should not throw exception when setting tokens for AM container.");
        }
    }