Finished writing tests. Some minor refactoring.