Improve test coverage, simplify code.