Cleanup and updated unit tests.