TestMultiTableInputFormatBase.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.hadoop.hbase.mapreduce;

import static org.apache.hadoop.hbase.client.Scan.SCAN_ATTRIBUTES_TABLE_NAME;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseClassTestRule;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.HConstants;
import org.apache.hadoop.hbase.HRegionLocation;
import org.apache.hadoop.hbase.ServerName;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Admin;
import org.apache.hadoop.hbase.client.AsyncConnection;
import org.apache.hadoop.hbase.client.BufferedMutator;
import org.apache.hadoop.hbase.client.BufferedMutatorParams;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionUtils;
import org.apache.hadoop.hbase.client.RegionInfoBuilder;
import org.apache.hadoop.hbase.client.RegionLocator;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.client.Scan;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.client.TableBuilder;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.security.User;
import org.apache.hadoop.hbase.testclassification.SmallTests;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.hbase.util.Pair;
import org.apache.hadoop.mapreduce.InputSplit;
import org.apache.hadoop.mapreduce.JobContext;
import org.apache.hadoop.mapreduce.RecordReader;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.junit.Assert;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.TestName;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;

/**
 * Tests of MultiTableInputFormatBase.
 */
@Category({SmallTests.class})
public class TestMultiTableInputFormatBase {

  @ClassRule
  public static final HBaseClassTestRule CLASS_RULE =
      HBaseClassTestRule.forClass(TestMultiTableInputFormatBase.class);

  @Rule public final TestName name = new TestName();

  /**
   * Test getSplits only puts up one Connection.
   * In past it has put up many Connections. Each Connection setup comes with a fresh new cache
   * so we have to do fresh hit on hbase:meta. Should only do one Connection when doing getSplits
   * even if a MultiTableInputFormat.
   * @throws IOException
   */
  @Test
  public void testMRSplitsConnectionCount() throws IOException {
    // Make instance of MTIFB.
    MultiTableInputFormatBase mtif = new MultiTableInputFormatBase() {
      @Override
      public RecordReader<ImmutableBytesWritable, Result> createRecordReader(InputSplit split,
          TaskAttemptContext context)
      throws IOException, InterruptedException {
        return super.createRecordReader(split, context);
      }
    };
    // Pass it a mocked JobContext. Make the JC return our Configuration.
    // Load the Configuration so it returns our special Connection so we can interpolate
    // canned responses.
    JobContext mockedJobContext = Mockito.mock(JobContext.class);
    Configuration c = HBaseConfiguration.create();
    c.set(ConnectionUtils.HBASE_CLIENT_CONNECTION_IMPL, MRSplitsConnection.class.getName());
    Mockito.when(mockedJobContext.getConfiguration()).thenReturn(c);
    // Invent a bunch of scans. Have each Scan go against a different table so a good spread.
    List<Scan> scans = new ArrayList<>();
    for (int i = 0; i < 10; i++) {
      Scan scan = new Scan();
      String tableName = this.name.getMethodName() + i;
      scan.setAttribute(SCAN_ATTRIBUTES_TABLE_NAME, Bytes.toBytes(tableName));
      scans.add(scan);
    }
    mtif.setScans(scans);
    // Get splits. Assert that that more than one.
    List<InputSplit> splits = mtif.getSplits(mockedJobContext);
    Assert.assertTrue(splits.size() > 0);
    // Assert only one Connection was made (see the static counter we have in the mocked
    // Connection MRSplitsConnection Constructor.
    Assert.assertEquals(1, MRSplitsConnection.creations.get());
  }

  /**
   * Connection to use above in Test.
   */
  public static class MRSplitsConnection implements Connection {
    private final Configuration configuration;
    static final AtomicInteger creations = new AtomicInteger(0);

    MRSplitsConnection (Configuration conf, ExecutorService pool, User user) throws IOException {
      this.configuration = conf;
      creations.incrementAndGet();
    }

    @Override
    public void abort(String why, Throwable e) {

    }

    @Override
    public boolean isAborted() {
      return false;
    }

    @Override
    public Configuration getConfiguration() {
      return this.configuration;
    }

    @Override
    public BufferedMutator getBufferedMutator(TableName tableName) throws IOException {
      return null;
    }

    @Override
    public BufferedMutator getBufferedMutator(BufferedMutatorParams params) throws IOException {
      return null;
    }

    @Override
    public RegionLocator getRegionLocator(final TableName tableName) throws IOException {
      // Make up array of start keys. We start off w/ empty byte array.
      final byte [][] startKeys = new byte [][] {HConstants.EMPTY_BYTE_ARRAY,
          Bytes.toBytes("aaaa"), Bytes.toBytes("bbb"),
          Bytes.toBytes("ccc"), Bytes.toBytes("ddd"), Bytes.toBytes("eee"),
          Bytes.toBytes("fff"), Bytes.toBytes("ggg"), Bytes.toBytes("hhh"),
          Bytes.toBytes("iii"), Bytes.toBytes("lll"), Bytes.toBytes("mmm"),
          Bytes.toBytes("nnn"), Bytes.toBytes("ooo"), Bytes.toBytes("ppp"),
          Bytes.toBytes("qqq"), Bytes.toBytes("rrr"), Bytes.toBytes("sss"),
          Bytes.toBytes("ttt"), Bytes.toBytes("uuu"), Bytes.toBytes("vvv"),
          Bytes.toBytes("zzz")};
      // Make an array of end keys. We end with the empty byte array.
      final byte [][] endKeys = new byte[][] {
          Bytes.toBytes("aaaa"), Bytes.toBytes("bbb"),
          Bytes.toBytes("ccc"), Bytes.toBytes("ddd"), Bytes.toBytes("eee"),
          Bytes.toBytes("fff"), Bytes.toBytes("ggg"), Bytes.toBytes("hhh"),
          Bytes.toBytes("iii"), Bytes.toBytes("lll"), Bytes.toBytes("mmm"),
          Bytes.toBytes("nnn"), Bytes.toBytes("ooo"), Bytes.toBytes("ppp"),
          Bytes.toBytes("qqq"), Bytes.toBytes("rrr"), Bytes.toBytes("sss"),
          Bytes.toBytes("ttt"), Bytes.toBytes("uuu"), Bytes.toBytes("vvv"),
          Bytes.toBytes("zzz"),
          HConstants.EMPTY_BYTE_ARRAY};
      // Now make a map of start keys to HRegionLocations. Let the server namber derive from
      // the start key.
      final Map<byte [], HRegionLocation> map =
          new TreeMap<byte [], HRegionLocation>(Bytes.BYTES_COMPARATOR);
      for (byte [] startKey: startKeys) {
        HRegionLocation hrl = new HRegionLocation(
            RegionInfoBuilder.newBuilder(tableName).setStartKey(startKey).build(),
            ServerName.valueOf(Bytes.toString(startKey), 0, 0));
        map.put(startKey, hrl);
      }
      // Get a list of the locations.
      final List<HRegionLocation> locations = new ArrayList<HRegionLocation>(map.values());
      // Now make a RegionLocator mock backed by the abpve map and list of locations.
      RegionLocator mockedRegionLocator = Mockito.mock(RegionLocator.class);
      Mockito.when(mockedRegionLocator.getRegionLocation(Mockito.any(byte [].class),
            Mockito.anyBoolean())).
          thenAnswer(new Answer<HRegionLocation>() {
            @Override
            public HRegionLocation answer(InvocationOnMock invocationOnMock) throws Throwable {
              Object [] args = invocationOnMock.getArguments();
              byte [] key = (byte [])args[0];
              return map.get(key);
            }
          });
      Mockito.when(mockedRegionLocator.getAllRegionLocations()).thenReturn(locations);
      Mockito.when(mockedRegionLocator.getStartEndKeys()).
          thenReturn(new Pair<byte [][], byte[][]>(startKeys, endKeys));
      Mockito.when(mockedRegionLocator.getName()).thenReturn(tableName);
      return mockedRegionLocator;
    }

    @Override
    public Admin getAdmin() throws IOException {
      Admin admin = Mockito.mock(Admin.class);
      Mockito.when(admin.getConfiguration()).thenReturn(getConfiguration());
      return admin;
    }

    @Override
    public Table getTable(TableName tableName) throws IOException {
      Table table = Mockito.mock(Table.class);
      Mockito.when(table.getName()).thenReturn(tableName);
      return table;
    }

    @Override
    public void close() throws IOException {

    }

    @Override
    public boolean isClosed() {
      return false;
    }

    @Override
    public TableBuilder getTableBuilder(TableName tableName, ExecutorService pool) {
      return Mockito.mock(TableBuilder.class);
    }

    @Override
    public void clearRegionLocationCache() {
    }

    @Override
    public AsyncConnection toAsyncConnection() {
      return null;
    }

    @Override
    public String getClusterId() {
      return null;
    }

  }
}