Fix unit test
Merge pull request #317 from tribal-tec/master

Fix unit test