Fix issues found writing tests