Examples with HBaseConnection com.amazonaws.athena.connectors.hbase.connection.HBaseConnection used on opensource projects

Example 1 with HBaseConnection

use of com.amazonaws.athena.connectors.hbase.connection.HBaseConnection in project aws-athena-query-federation by awslabs.

the class HbaseSchemaUtilsTest method inferSchema.

@Test
public void inferSchema() throws IOException {
    int numToScan = 4;
    TableName tableName = new TableName("schema", "table");
    List<Result> results = TestUtils.makeResults();
    HBaseConnection mockConnection = mock(HBaseConnection.class);
    ResultScanner mockScanner = mock(ResultScanner.class);
    when(mockScanner.iterator()).thenReturn(results.iterator());
    when(mockConnection.scanTable(anyObject(), any(Scan.class), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        ResultProcessor processor = (ResultProcessor) invocationOnMock.getArguments()[2];
        return processor.scan(mockScanner);
    });
    Schema schema = HbaseSchemaUtils.inferSchema(mockConnection, tableName, numToScan);
    Map<String, Types.MinorType> actualFields = new HashMap<>();
    schema.getFields().stream().forEach(next -> actualFields.put(next.getName(), Types.getMinorTypeForArrowType(next.getType())));
    Map<String, Types.MinorType> expectedFields = new HashMap<>();
    TestUtils.makeSchema().build().getFields().stream().forEach(next -> expectedFields.put(next.getName(), Types.getMinorTypeForArrowType(next.getType())));
    for (Map.Entry<String, Types.MinorType> nextExpected : expectedFields.entrySet()) {
        assertNotNull(actualFields.get(nextExpected.getKey()));
        assertEquals(nextExpected.getKey(), nextExpected.getValue(), actualFields.get(nextExpected.getKey()));
    }
    assertEquals(expectedFields.size(), actualFields.size());
    verify(mockConnection, times(1)).scanTable(anyObject(), any(Scan.class), any(ResultProcessor.class));
    verify(mockScanner, times(1)).iterator();
}