tests: Fix rebalancing test

Description in commit