Fixing unit tests (#412)