Clean up code, bring unit tests to 100.