Simplify code, improve test coverage.