Proper fix / More edge case in the test