Fix corresponding unit test

Had already pushed the change, so I couldn't amend it!
Merge pull request #4 from jtoberon/master

Fix corresponding unit test