Updates to improve test coverage