Added error handling in testing