Unit tests fix. At the same time fixed some errors :)