Improve test coverage and some code cleanup.