Fixing unit tests

Hey Rob, 

This fixed the unit test for me, please confirm and merge.

Thanks