Cleaning up and adding tests.