Fixing test cases, have RebalanceController use ProtoBuf