Test cleanup

@hlbraddock review? :smile_cat: