Refactor tests to have QuickCheck shrink