Cleaning up and adding null tests