}
@Test
public void testAssignMigrations() {
servers.clear();
servers.put(new TServerInstance(AddressUtil.parseAddress("127.0.0.1", 1234), "a"), new FakeTServer());
servers.put(new TServerInstance(AddressUtil.parseAddress("127.0.0.1", 1235), "b"), new FakeTServer());
servers.put(new TServerInstance(AddressUtil.parseAddress("127.0.0.1", 1236), "c"), new FakeTServer());
List<KeyExtent> metadataTable = new ArrayList<KeyExtent>();
String table = "t1";
metadataTable.add(makeExtent(table, null, null));
table = "t2";
metadataTable.add(makeExtent(table, "a", null));
metadataTable.add(makeExtent(table, null, "a"));
table = "t3";
metadataTable.add(makeExtent(table, "a", null));
metadataTable.add(makeExtent(table, "b", "a"));
metadataTable.add(makeExtent(table, "c", "b"));
metadataTable.add(makeExtent(table, "d", "c"));
metadataTable.add(makeExtent(table, "e", "d"));
metadataTable.add(makeExtent(table, null, "e"));
Collections.sort(metadataTable);
TestDefaultLoadBalancer balancer = new TestDefaultLoadBalancer();
SortedMap<TServerInstance,TabletServerStatus> current = new TreeMap<TServerInstance,TabletServerStatus>();
for (Entry<TServerInstance,FakeTServer> entry : servers.entrySet()) {
current.put(entry.getKey(), entry.getValue().getStatus(entry.getKey()));
}
assignTablets(metadataTable, servers, current, balancer);
// Verify that the counts on the tables are correct
Map<String,Integer> expectedCounts = new HashMap<String,Integer>();
expectedCounts.put("t1", 1);
expectedCounts.put("t2", 1);
expectedCounts.put("t3", 2);
checkBalance(metadataTable, servers, expectedCounts);
// Rebalance once
for (Entry<TServerInstance,FakeTServer> entry : servers.entrySet()) {
current.put(entry.getKey(), entry.getValue().getStatus(entry.getKey()));
}
// Nothing should happen, we are balanced
ArrayList<TabletMigration> out = new ArrayList<TabletMigration>();
balancer.getMigrations(current, out);
assertEquals(out.size(), 0);
// Take down a tabletServer
TServerInstance first = current.keySet().iterator().next();
current.remove(first);
FakeTServer remove = servers.remove(first);
// reassign offline extents
assignTablets(remove.extents, servers, current, balancer);