Refactored tests a little.