Update randomized test to use ScalaTest's assume()

Author: JoshRosen

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/CodeGenerationSuite.scala

@@ -60,22 +60,23 @@ class CodeGenerationSuite extends SparkFunSuite {
       val rowType = StructType(
         StructField("a", dataType, nullable = true) ::
         StructField("b", dataType, nullable = true) :: Nil)
+      val maybeDataGenerator = RandomDataGenerator.forType(rowType, nullable = false)
+      assume(maybeDataGenerator.isDefined)
+      val randGenerator = maybeDataGenerator.get
       val toCatalyst = CatalystTypeConverters.createToCatalystConverter(rowType)
-      RandomDataGenerator.forType(rowType, nullable = false).foreach { randGenerator =>
-        for (_ <- 1 to 50) {
-          val a = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
-          val b = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
-          withClue(s"a = $a, b = $b") {
-            assert(genOrdering.compare(a, a) === 0)
-            assert(genOrdering.compare(b, b) === 0)
-            assert(rowOrdering.compare(a, a) === 0)
-            assert(rowOrdering.compare(b, b) === 0)
-            assert(signum(genOrdering.compare(a, b)) === -1 * signum(genOrdering.compare(b, a)))
-            assert(signum(rowOrdering.compare(a, b)) === -1 * signum(rowOrdering.compare(b, a)))
-            assert(
-              signum(rowOrdering.compare(a, b)) === signum(genOrdering.compare(a, b)),
-              "Generated and non-generated orderings should agree")
-          }
+      for (_ <- 1 to 50) {
+        val a = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
+        val b = toCatalyst(randGenerator()).asInstanceOf[InternalRow]
+        withClue(s"a = $a, b = $b") {
+          assert(genOrdering.compare(a, a) === 0)
+          assert(genOrdering.compare(b, b) === 0)
+          assert(rowOrdering.compare(a, a) === 0)
+          assert(rowOrdering.compare(b, b) === 0)
+          assert(signum(genOrdering.compare(a, b)) === -1 * signum(genOrdering.compare(b, a)))
+          assert(signum(rowOrdering.compare(a, b)) === -1 * signum(rowOrdering.compare(b, a)))
+          assert(
+            signum(rowOrdering.compare(a, b)) === signum(genOrdering.compare(a, b)),
+            "Generated and non-generated orderings should agree")
         }
       }
     }