Fix broken unit test

I missed something in PR #9, fixing.
Merge pull request #10 from beav/unittest

Fix broken unit test