Clean/tests

cc @lancejpollard