Some code cleanup and tests.