Fix index out of bounds, better test case