Fixing test classes