示例#1
0
  /**
   * Tests whether data in valuesIterator matches with sorted input data set.
   *
   * <p>Returns a list of value counts for each key.
   *
   * @param valuesIterator
   * @return List
   * @throws IOException
   */
  private List<Integer> verifyIteratorData(ValuesIterator valuesIterator) throws IOException {
    boolean result = true;

    ArrayList<Integer> sequence = new ArrayList<Integer>();

    // sort original data based on comparator
    ListMultimap<Writable, Writable> sortedMap =
        new ImmutableListMultimap.Builder<Writable, Writable>()
            .orderKeysBy(this.correctComparator)
            .putAll(originalData)
            .build();

    Set<Map.Entry<Writable, Writable>> oriKeySet = Sets.newSet();
    oriKeySet.addAll(sortedMap.entries());

    // Iterate through sorted data and valuesIterator for verification
    for (Map.Entry<Writable, Writable> entry : oriKeySet) {

      assertTrue(valuesIterator.moveToNext());

      Writable oriKey = entry.getKey();
      // Verify if the key and the original key are same
      if (!oriKey.equals((Writable) valuesIterator.getKey())) {
        result = false;
        break;
      }

      int valueCount = 0;
      // Verify values
      Iterator<Writable> vItr = valuesIterator.getValues().iterator();
      for (Writable val : sortedMap.get(oriKey)) {
        assertTrue(vItr.hasNext());

        // Verify if the values are same
        if (!val.equals((Writable) vItr.next())) {
          result = false;
          break;
        }

        valueCount++;
      }
      sequence.add(valueCount);
      assertTrue("At least 1 value per key", valueCount > 0);
    }
    if (expectedTestResult) {
      assertTrue(result);

      assertFalse(valuesIterator.moveToNext());
      getNextFromFinishedIterator(valuesIterator);
    } else {
      while (valuesIterator.moveToNext()) {
        // iterate through all keys
      }
      getNextFromFinishedIterator(valuesIterator);
      assertFalse(result);
    }

    return sequence;
  }
  @Test
  public void testGrantedAuthoritiesMapper() throws Exception {
    SimpleAuthorityMapper grantedAuthorityMapper = new SimpleAuthorityMapper();
    grantedAuthorityMapper.setPrefix("ROLE_");
    grantedAuthorityMapper.setConvertToUpperCase(true);
    provider.setGrantedAuthoritiesMapper(grantedAuthorityMapper);

    Authentication result = provider.authenticate(token);
    assertEquals(
        Sets.newSet("ROLE_USER", "ROLE_ADMIN"),
        AuthorityUtils.authorityListToSet(result.getAuthorities()));
  }
/** Keycloak authentication provider tests. */
public class KeycloakAuthenticationProviderTest {
  private KeycloakAuthenticationProvider provider = new KeycloakAuthenticationProvider();
  private KeycloakAuthenticationToken token;
  private Set<String> roles = Sets.newSet("user", "admin");

  @Before
  public void setUp() throws Exception {
    Principal principal = mock(Principal.class);
    RefreshableKeycloakSecurityContext securityContext =
        mock(RefreshableKeycloakSecurityContext.class);
    KeycloakAccount account = new SimpleKeycloakAccount(principal, roles, securityContext);

    token = new KeycloakAuthenticationToken(account);
  }

  @Test
  public void testAuthenticate() throws Exception {
    Authentication result = provider.authenticate(token);
    assertNotNull(result);
    assertEquals(roles, AuthorityUtils.authorityListToSet(result.getAuthorities()));
    assertTrue(result.isAuthenticated());
    assertNotNull(result.getPrincipal());
    assertNotNull(result.getCredentials());
    assertNotNull(result.getDetails());
  }

  @Test
  public void testSupports() throws Exception {
    assertTrue(provider.supports(KeycloakAuthenticationToken.class));
    assertFalse(provider.supports(PreAuthenticatedAuthenticationToken.class));
  }

  @Test
  public void testGrantedAuthoritiesMapper() throws Exception {
    SimpleAuthorityMapper grantedAuthorityMapper = new SimpleAuthorityMapper();
    grantedAuthorityMapper.setPrefix("ROLE_");
    grantedAuthorityMapper.setConvertToUpperCase(true);
    provider.setGrantedAuthoritiesMapper(grantedAuthorityMapper);

    Authentication result = provider.authenticate(token);
    assertEquals(
        Sets.newSet("ROLE_USER", "ROLE_ADMIN"),
        AuthorityUtils.authorityListToSet(result.getAuthorities()));
  }
}
  @Test(timeout = 5000)
  public void testDAGClientHandler() throws TezException {

    TezDAGID mockTezDAGId = mock(TezDAGID.class);
    when(mockTezDAGId.getId()).thenReturn(1);
    when(mockTezDAGId.toString()).thenReturn("dag_9999_0001_1");

    DAG mockDAG = mock(DAG.class);
    when(mockDAG.getID()).thenReturn(mockTezDAGId);
    DAGStatusBuilder mockDagStatusBuilder = mock(DAGStatusBuilder.class);
    when(mockDAG.getDAGStatus(anySetOf(StatusGetOpts.class))).thenReturn(mockDagStatusBuilder);
    VertexStatusBuilder mockVertexStatusBuilder = mock(VertexStatusBuilder.class);
    when(mockDAG.getVertexStatus(anyString(), anySetOf(StatusGetOpts.class)))
        .thenReturn(mockVertexStatusBuilder);

    DAGAppMaster mockDagAM = mock(DAGAppMaster.class);
    when(mockDagAM.getState()).thenReturn(DAGAppMasterState.RUNNING);
    AppContext mockAppContext = mock(AppContext.class);
    when(mockDagAM.getContext()).thenReturn(mockAppContext);
    when(mockDagAM.getContext().getCurrentDAG()).thenReturn(mockDAG);

    DAGClientHandler dagClientHandler = new DAGClientHandler(mockDagAM);

    // getAllDAGs()
    assertEquals(1, dagClientHandler.getAllDAGs().size());
    assertEquals("dag_9999_0001_1", dagClientHandler.getAllDAGs().get(0));

    // getDAGStatus
    try {
      dagClientHandler.getDAGStatus("dag_9999_0001_2", Sets.newSet(StatusGetOpts.GET_COUNTERS));
      fail("should not come here");
    } catch (TezException e) {
      assertTrue(e.getMessage().contains("Unknown dagId"));
    }
    DAGStatus dagStatus =
        dagClientHandler.getDAGStatus("dag_9999_0001_1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
    assertEquals(mockDagStatusBuilder, dagStatus);

    // getVertexStatus
    try {
      dagClientHandler.getVertexStatus(
          "dag_9999_0001_2", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
      fail("should not come here");
    } catch (TezException e) {
      assertTrue(e.getMessage().contains("Unknown dagId"));
    }
    VertexStatus vertexStatus =
        dagClientHandler.getVertexStatus(
            "dag_9999_0001_1", "v1", Sets.newSet(StatusGetOpts.GET_COUNTERS));
    assertEquals(mockVertexStatusBuilder, vertexStatus);

    // getTezAppMasterStatus
    when(mockDagAM.isSession()).thenReturn(false);

    assertEquals(TezAppMasterStatus.RUNNING, dagClientHandler.getTezAppMasterStatus());

    when(mockDagAM.isSession()).thenReturn(true);
    when(mockDagAM.getState()).thenReturn(DAGAppMasterState.INITED);
    assertEquals(TezAppMasterStatus.INITIALIZING, dagClientHandler.getTezAppMasterStatus());
    when(mockDagAM.getState()).thenReturn(DAGAppMasterState.ERROR);
    assertEquals(TezAppMasterStatus.SHUTDOWN, dagClientHandler.getTezAppMasterStatus());

    // tryKillDAG
    try {
      dagClientHandler.tryKillDAG("dag_9999_0001_2");
      fail("should not come here");
    } catch (TezException e) {
      assertTrue(e.getMessage().contains("Unknown dagId"));
    }
    dagClientHandler.tryKillDAG("dag_9999_0001_1");
    ArgumentCaptor<DAG> eventCaptor = ArgumentCaptor.forClass(DAG.class);
    verify(mockDagAM, times(1))
        .tryKillDAG(eventCaptor.capture(), contains("Sending client kill from"));
    assertEquals(1, eventCaptor.getAllValues().size());
    assertTrue(eventCaptor.getAllValues().get(0) instanceof DAG);
    assertEquals("dag_9999_0001_1", ((DAG) eventCaptor.getAllValues().get(0)).getID().toString());

    // submitDAG
    DAGPlan dagPlan = DAGPlan.getDefaultInstance();
    Map<String, LocalResource> localResources = new HashMap<String, LocalResource>();
    dagClientHandler.submitDAG(dagPlan, localResources);
    verify(mockDagAM).submitDAGToAppMaster(dagPlan, localResources);

    // shutdown
    dagClientHandler.shutdownAM();
    verify(mockDagAM).shutdownTezAM(contains("Received message to shutdown AM from"));
  }