@Test(timeout = 20000) public void testGetHSDelegationToken() throws Exception { try { Configuration conf = new Configuration(); // Setup mock service InetSocketAddress mockRmAddress = new InetSocketAddress("localhost", 4444); Text rmTokenSevice = SecurityUtil.buildTokenService(mockRmAddress); InetSocketAddress mockHsAddress = new InetSocketAddress("localhost", 9200); Text hsTokenSevice = SecurityUtil.buildTokenService(mockHsAddress); // Setup mock rm token RMDelegationTokenIdentifier tokenIdentifier = new RMDelegationTokenIdentifier(new Text("owner"), new Text("renewer"), new Text("real")); Token<RMDelegationTokenIdentifier> token = new Token<RMDelegationTokenIdentifier>( new byte[0], new byte[0], tokenIdentifier.getKind(), rmTokenSevice); token.setKind(RMDelegationTokenIdentifier.KIND_NAME); // Setup mock history token DelegationToken historyToken = BuilderUtils.newDelegationToken( new byte[0], MRDelegationTokenIdentifier.KIND_NAME.toString(), new byte[0], hsTokenSevice.toString()); GetDelegationTokenResponse getDtResponse = Records.newRecord(GetDelegationTokenResponse.class); getDtResponse.setDelegationToken(historyToken); // mock services MRClientProtocol mockHsProxy = mock(MRClientProtocol.class); doReturn(mockHsAddress).when(mockHsProxy).getConnectAddress(); doReturn(getDtResponse) .when(mockHsProxy) .getDelegationToken(any(GetDelegationTokenRequest.class)); ResourceMgrDelegate rmDelegate = mock(ResourceMgrDelegate.class); doReturn(mockRmAddress).when(rmDelegate).getConnectAddress(); ClientCache clientCache = mock(ClientCache.class); doReturn(mockHsProxy).when(clientCache).getInitializedHSProxy(); Credentials creds = new Credentials(); YARNRunner yarnRunner = new YARNRunner(conf, rmDelegate, clientCache); // No HS token if no RM token yarnRunner.addHistoyToken(creds); verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class)); // No HS token if RM token, but secirity disabled. creds.addToken(new Text("rmdt"), token); yarnRunner.addHistoyToken(creds); verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class)); conf.set(CommonConfigurationKeys.HADOOP_SECURITY_AUTHENTICATION, "kerberos"); UserGroupInformation.setConfiguration(conf); creds = new Credentials(); // No HS token if no RM token, security enabled yarnRunner.addHistoyToken(creds); verify(mockHsProxy, times(0)).getDelegationToken(any(GetDelegationTokenRequest.class)); // HS token if RM token present, security enabled creds.addToken(new Text("rmdt"), token); yarnRunner.addHistoyToken(creds); verify(mockHsProxy, times(1)).getDelegationToken(any(GetDelegationTokenRequest.class)); // No additional call to get HS token if RM and HS token present yarnRunner.addHistoyToken(creds); verify(mockHsProxy, times(1)).getDelegationToken(any(GetDelegationTokenRequest.class)); } finally { // Back to defaults. UserGroupInformation.setConfiguration(new Configuration()); } }