Fix broken test.
Merge pull request #152 from square/dimitris/fix-test

Fix broken test.