Fixed and improved a unit test