Fixing the broken test