Add rebalance-related tests