Testing operations that cause a shuffle in Apache Spark