clean up and updating tests