1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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
39
40 public class HConnectionTestingUtility {
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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 ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
67 }
68 return connection;
69 }
70 }
71
72
73
74
75 private static void mockRegionLocator(final HConnectionImplementation connection) {
76 try {
77 Mockito.when(connection.getRegionLocator(Mockito.any(TableName.class))).thenAnswer(
78 new Answer<RegionLocator>() {
79 @Override
80 public RegionLocator answer(InvocationOnMock invocation) throws Throwable {
81 TableName tableName = (TableName) invocation.getArguments()[0];
82 return new HRegionLocator(tableName, connection);
83 }
84 });
85 } catch (IOException e) {
86 }
87 }
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118 public static ClusterConnection getMockedConnectionAndDecorate(final Configuration conf,
119 final AdminProtos.AdminService.BlockingInterface admin,
120 final ClientProtos.ClientService.BlockingInterface client,
121 final ServerName sn, final HRegionInfo hri)
122 throws IOException {
123 HConnectionImplementation c = Mockito.mock(HConnectionImplementation.class);
124 Mockito.when(c.getConfiguration()).thenReturn(conf);
125 ConnectionManager.CONNECTION_INSTANCES.put(new HConnectionKey(conf), c);
126 Mockito.doNothing().when(c).close();
127
128 final HRegionLocation loc = new HRegionLocation(hri, sn);
129 mockRegionLocator(c);
130 Mockito.when(c.getRegionLocation((TableName) Mockito.any(),
131 (byte[]) Mockito.any(), Mockito.anyBoolean())).
132 thenReturn(loc);
133 Mockito.when(c.locateRegion((TableName) Mockito.any(), (byte[]) Mockito.any())).
134 thenReturn(loc);
135 Mockito.when(c.locateRegion((TableName) Mockito.any(), (byte[]) Mockito.any(),
136 Mockito.anyBoolean(), Mockito.anyBoolean(), Mockito.anyInt()))
137 .thenReturn(new RegionLocations(loc));
138 if (admin != null) {
139
140 Mockito.when(c.getAdmin(Mockito.any(ServerName.class))).
141 thenReturn(admin);
142 }
143 if (client != null) {
144
145 Mockito.when(c.getClient(Mockito.any(ServerName.class))).
146 thenReturn(client);
147 }
148 NonceGenerator ng = Mockito.mock(NonceGenerator.class);
149 Mockito.when(c.getNonceGenerator()).thenReturn(ng);
150 Mockito.when(c.getAsyncProcess()).thenReturn(
151 new AsyncProcess(c, conf, null, RpcRetryingCallerFactory.instantiate(conf), false,
152 RpcControllerFactory.instantiate(conf)));
153 Mockito.doNothing().when(c).incCount();
154 Mockito.doNothing().when(c).decCount();
155 Mockito.when(c.getNewRpcRetryingCallerFactory(conf)).thenReturn(
156 RpcRetryingCallerFactory.instantiate(conf,
157 RetryingCallerInterceptorFactory.NO_OP_INTERCEPTOR, null));
158 HTableInterface t = Mockito.mock(HTableInterface.class);
159 Mockito.when(c.getTable((TableName)Mockito.any())).thenReturn(t);
160 ResultScanner rs = Mockito.mock(ResultScanner.class);
161 Mockito.when(t.getScanner((Scan)Mockito.any())).thenReturn(rs);
162 return c;
163 }
164
165
166
167
168
169
170
171
172
173
174
175
176
177 public static ClusterConnection getSpiedConnection(final Configuration conf)
178 throws IOException {
179 HConnectionKey connectionKey = new HConnectionKey(conf);
180 synchronized (ConnectionManager.CONNECTION_INSTANCES) {
181 HConnectionImplementation connection =
182 ConnectionManager.CONNECTION_INSTANCES.get(connectionKey);
183 if (connection == null) {
184 connection = Mockito.spy(new HConnectionImplementation(conf, true));
185 ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
186 }
187 return connection;
188 }
189 }
190
191 public static ClusterConnection getSpiedClusterConnection(final Configuration conf)
192 throws IOException {
193 HConnectionKey connectionKey = new HConnectionKey(conf);
194 synchronized (ConnectionManager.CONNECTION_INSTANCES) {
195 HConnectionImplementation connection =
196 ConnectionManager.CONNECTION_INSTANCES.get(connectionKey);
197 if (connection == null) {
198 connection = Mockito.spy(new HConnectionImplementation(conf, true));
199 ConnectionManager.CONNECTION_INSTANCES.put(connectionKey, connection);
200 }
201 return connection;
202 }
203 }
204
205
206
207
208 public static int getConnectionCount() {
209 synchronized (ConnectionManager.CONNECTION_INSTANCES) {
210 return ConnectionManager.CONNECTION_INSTANCES.size();
211 }
212 }
213 }