Clean up code duplication in tests