Small tidying of previous commit

Very minor improvement of unit test and fix indentation issue.