Improved unit tests by using MagicMock and testing the task return value