fixing some edge cases in the tests