MAGE-726: Fix tests
- added return statements removed by mistake