Cleaning code up somemore based on test output