Quick wee refactor
Tidied up unit tests