Fixing test code