More test coverage and fixes