Updated tests to cover more cases.