Fix return code checking