cleaning up and improving tests