Fix UTs
Merge pull request #320 from ovatsus/master

Fix UTs