Fixing test cleanup