Tweaks test code

Fixes #31: Tweaks test code