Add tests for error handling