Improve warning check for ufunc test