Better error handling and test coverage