View Javadoc

1   /**
2    * Licensed to the Apache Software Foundation (ASF) under one
3    * or more contributor license agreements.  See the NOTICE file
4    * distributed with this work for additional information
5    * regarding copyright ownership.  The ASF licenses this file
6    * to you under the Apache License, Version 2.0 (the
7    * "License"); you may not use this file except in compliance
8    * with the License.  You may obtain a copy of the License at
9    *
10   *     http://www.apache.org/licenses/LICENSE-2.0
11   *
12   * Unless required by applicable law or agreed to in writing, software
13   * distributed under the License is distributed on an "AS IS" BASIS,
14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15   * See the License for the specific language governing permissions and
16   * limitations under the License.
17   */
18  package org.apache.hadoop.hbase.client;
19  
20  import java.io.IOException;
21  
22  import org.apache.hadoop.conf.Configuration;
23  import org.apache.hadoop.hbase.RegionLocations;
24  import org.apache.hadoop.hbase.TableName;
25  import org.apache.hadoop.hbase.HRegionInfo;
26  import org.apache.hadoop.hbase.HRegionLocation;
27  import org.apache.hadoop.hbase.ServerName;
28  import org.apache.hadoop.hbase.ZooKeeperConnectionException;
29  import org.apache.hadoop.hbase.protobuf.generated.AdminProtos;
30  import org.apache.hadoop.hbase.protobuf.generated.ClientProtos;
31  import org.apache.hadoop.hbase.client.ConnectionManager.HConnectionImplementation;
32  import org.apache.hadoop.hbase.ipc.RpcControllerFactory;
33  import org.mockito.Mockito;
34  import org.mockito.invocation.InvocationOnMock;
35  import org.mockito.stubbing.Answer;
36  
37  /**
38   * {@link ClusterConnection} testing utility.
39   */
40  public class HConnectionTestingUtility {
41    /*
42     * Not part of {@link HBaseTestingUtility} because this class is not
43     * in same package as {@link HConnection}.  Would have to reveal ugly
44     * {@link HConnectionManager} innards to HBaseTestingUtility to give it access.
45     */
46    /**
47     * Get a Mocked {@link HConnection} that goes with the passed <code>conf</code>
48     * configuration instance.  Minimally the mock will return
49     * <code>conf</conf> when {@link ClusterConnection#getConfiguration()} is invoked.
50     * Be sure to shutdown the connection when done by calling
51     * {@link HConnectionManager#deleteConnection(Configuration)} else it
52     * will stick around; this is probably not what you want.
53     * @param conf configuration
54     * @return HConnection object for <code>conf</code>
55     * @throws ZooKeeperConnectionException
56     */
57    public static ClusterConnection getMockedConnection(final Configuration conf)
58    throws ZooKeeperConnectionException {
59      HConnectionKey connectionKey = new HConnectionKey(conf);
60      synchronized (ConnectionManager.CONNECTION_INSTANCES) {
61        HConnectionImplementation connection =
62            ConnectionManager.CONNECTION_INSTANCES.get(connectionKey);
63        if (connection == null) {
64          connection = Mockito.mock(HConnectionImplementation.class);
65          Mockito.when(connection.getConfiguration()).thenReturn(conf);
66          Mockito.when(connection.getRpcControllerFactory()).thenReturn(
67          Mockito.mock(RpcControllerFactory.class));
68          // we need a real retrying caller
69          RpcRetryingCallerFactory callerFactory = new RpcRetryingCallerFactory(conf);
70          Mockito.when(connection.getRpcRetryingCallerFactory()).thenReturn(callerFactory);
71          ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
72        }
73        return connection;
74      }
75    }
76  
77    /**
78     * @param connection
79     */
80    private static void mockRegionLocator(final HConnectionImplementation connection) {
81      try {
82        Mockito.when(connection.getRegionLocator(Mockito.any(TableName.class))).thenAnswer(
83            new Answer<RegionLocator>() {
84              @Override
85              public RegionLocator answer(InvocationOnMock invocation) throws Throwable {
86                TableName tableName = (TableName) invocation.getArguments()[0];
87                return new HRegionLocator(tableName, connection);
88              }
89            });
90      } catch (IOException e) {
91      }
92    }
93  
94    /**
95     * Calls {@link #getMockedConnection(Configuration)} and then mocks a few
96     * more of the popular {@link ClusterConnection} methods so they do 'normal'
97     * operation (see return doc below for list). Be sure to shutdown the
98     * connection when done by calling
99     * {@link HConnectionManager#deleteConnection(Configuration)} else it
100    * will stick around; this is probably not what you want.
101    *
102    * @param conf Configuration to use
103    * @param admin An AdminProtocol; can be null but is usually
104    * itself a mock.
105    * @param client A ClientProtocol; can be null but is usually
106    * itself a mock.
107    * @param sn ServerName to include in the region location returned by this
108    * <code>connection</code>
109    * @param hri HRegionInfo to include in the location returned when
110    * getRegionLocator is called on the mocked connection
111    * @return Mock up a connection that returns a {@link Configuration} when
112    * {@link ClusterConnection#getConfiguration()} is called, a 'location' when
113    * {@link ClusterConnection#getRegionLocation(org.apache.hadoop.hbase.TableName, byte[], boolean)}
114    * is called,
115    * and that returns the passed {@link AdminProtos.AdminService.BlockingInterface} instance when
116    * {@link ClusterConnection#getAdmin(ServerName)} is called, returns the passed
117    * {@link ClientProtos.ClientService.BlockingInterface} instance when
118    * {@link ClusterConnection#getClient(ServerName)} is called (Be sure to call
119    * {@link HConnectionManager#deleteConnection(Configuration)}
120    * when done with this mocked Connection.
121    * @throws IOException
122    */
123   public static ClusterConnection getMockedConnectionAndDecorate(final Configuration conf,
124       final AdminProtos.AdminService.BlockingInterface admin,
125       final ClientProtos.ClientService.BlockingInterface client,
126       final ServerName sn, final HRegionInfo hri)
127   throws IOException {
128     HConnectionImplementation c = Mockito.mock(HConnectionImplementation.class);
129     Mockito.when(c.getConfiguration()).thenReturn(conf);
130     ConnectionManager.CONNECTION_INSTANCES.put(new HConnectionKey(conf), c);
131     Mockito.doNothing().when(c).close();
132     // Make it so we return a particular location when asked.
133     final HRegionLocation loc = new HRegionLocation(hri, sn);
134     mockRegionLocator(c);
135     Mockito.when(c.getRegionLocation((TableName) Mockito.any(),
136         (byte[]) Mockito.any(), Mockito.anyBoolean())).
137       thenReturn(loc);
138     Mockito.when(c.locateRegion((TableName) Mockito.any(), (byte[]) Mockito.any())).
139       thenReturn(loc);
140     Mockito.when(c.locateRegion((TableName) Mockito.any(), (byte[]) Mockito.any(),
141         Mockito.anyBoolean(), Mockito.anyBoolean(),  Mockito.anyInt()))
142         .thenReturn(new RegionLocations(loc));
143     if (admin != null) {
144       // If a call to getAdmin, return this implementation.
145       Mockito.when(c.getAdmin(Mockito.any(ServerName.class))).
146         thenReturn(admin);
147     }
148     if (client != null) {
149       // If a call to getClient, return this client.
150       Mockito.when(c.getClient(Mockito.any(ServerName.class))).
151         thenReturn(client);
152     }
153     NonceGenerator ng = Mockito.mock(NonceGenerator.class);
154     Mockito.when(c.getNonceGenerator()).thenReturn(ng);
155     Mockito.when(c.getAsyncProcess()).thenReturn(
156       new AsyncProcess(c, conf, null, RpcRetryingCallerFactory.instantiate(conf), false,
157           RpcControllerFactory.instantiate(conf)));
158     Mockito.doNothing().when(c).incCount();
159     Mockito.doNothing().when(c).decCount();
160     Mockito.when(c.getNewRpcRetryingCallerFactory(conf)).thenReturn(
161         RpcRetryingCallerFactory.instantiate(conf,
162             RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, null));
163     Mockito.when(c.getRpcControllerFactory()).thenReturn(Mockito.mock(RpcControllerFactory.class));
164     HTableInterface t = Mockito.mock(HTableInterface.class);
165     Mockito.when(c.getTable((TableName)Mockito.any())).thenReturn(t);
166     ResultScanner rs = Mockito.mock(ResultScanner.class);
167     Mockito.when(t.getScanner((Scan)Mockito.any())).thenReturn(rs);
168     return c;
169   }
170 
171   /**
172    * Get a Mockito spied-upon {@link ClusterConnection} that goes with the passed
173    * <code>conf</code> configuration instance.
174    * Be sure to shutdown the connection when done by calling
175    * {@link HConnectionManager#deleteConnection(Configuration)} else it
176    * will stick around; this is probably not what you want.
177    * @param conf configuration
178    * @return HConnection object for <code>conf</code>
179    * @throws ZooKeeperConnectionException
180    * @see @link
181    * {http://mockito.googlecode.com/svn/branches/1.6/javadoc/org/mockito/Mockito.html#spy(T)}
182    */
183   public static ClusterConnection getSpiedConnection(final Configuration conf)
184   throws IOException {
185     HConnectionKey connectionKey = new HConnectionKey(conf);
186     synchronized (ConnectionManager.CONNECTION_INSTANCES) {
187       HConnectionImplementation connection =
188           ConnectionManager.CONNECTION_INSTANCES.get(connectionKey);
189       if (connection == null) {
190         connection = Mockito.spy(new HConnectionImplementation(conf, true));
191         ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
192       }
193       return connection;
194     }
195   }
196 
197   public static ClusterConnection getSpiedClusterConnection(final Configuration conf)
198   throws IOException {
199     HConnectionKey connectionKey = new HConnectionKey(conf);
200     synchronized (ConnectionManager.CONNECTION_INSTANCES) {
201       HConnectionImplementation connection =
202           ConnectionManager.CONNECTION_INSTANCES.get(connectionKey);
203       if (connection == null) {
204         connection = Mockito.spy(new HConnectionImplementation(conf, true));
205         ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
206       }
207       return connection;
208     }
209   }
210 
211   /**
212    * @return Count of extant connection instances
213    */
214   public static int getConnectionCount() {
215     synchronized (ConnectionManager.CONNECTION_INSTANCES) {
216       return ConnectionManager.CONNECTION_INSTANCES.size();
217     }
218   }
219 }