Fixing other issue with unit tests