Add some extra checks for error handling