[SPARK-25456][SQL][TEST] Fix PythonForeachWriterSuite

squito · squito · commit a6f37b0742d8 · 2018-09-18T16:33:37.000-05:00
PythonForeachWriterSuite was failing because RowQueue now needs to have a handle on a SparkEnv with a SerializerManager, so added a mock env with a serializer manager. Also fixed a typo in the `finally` that was hiding the real exception. Tested PythonForeachWriterSuite locally, full tests via jenkins. Closes #22452 from squito/SPARK-25456. Authored-by: Imran Rashid <irashid@cloudera.com> Signed-off-by: Imran Rashid <irashid@cloudera.com>
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonForeachWriterSuite.scala
@@ -19,17 +19,20 @@ package org.apache.spark.sql.execution.python
 
 import scala.collection.mutable.ArrayBuffer
 
+import org.mockito.Mockito.when
 import org.scalatest.concurrent.Eventually
+import org.scalatest.mockito.MockitoSugar
 import org.scalatest.time.SpanSugar._
 
 import org.apache.spark._
 import org.apache.spark.memory.{TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.serializer.{JavaSerializer, SerializerManager}
 import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, UnsafeProjection}
 import org.apache.spark.sql.execution.python.PythonForeachWriter.UnsafeRowBuffer
 import org.apache.spark.sql.types.{DataType, IntegerType}
 import org.apache.spark.util.Utils
 
-class PythonForeachWriterSuite extends SparkFunSuite with Eventually {
+class PythonForeachWriterSuite extends SparkFunSuite with Eventually with MockitoSugar {
 
   testWithBuffer("UnsafeRowBuffer: iterator blocks when no data is available") { b =>
     b.assertIteratorBlocked()
@@ -75,15 +78,20 @@ class PythonForeachWriterSuite extends SparkFunSuite with Eventually {
         tester = new BufferTester(memBytes, sleepPerRowReadMs)
         f(tester)
       } finally {
-        if (tester == null) tester.close()
+        if (tester != null) tester.close()
       }
     }
   }
 
 
   class BufferTester(memBytes: Long, sleepPerRowReadMs: Int) {
     private val buffer = {
-      val mem = new TestMemoryManager(new SparkConf())
+      val mockEnv = mock[SparkEnv]
+      val conf = new SparkConf()
+      val serializerManager = new SerializerManager(new JavaSerializer(conf), conf, None)
+      when(mockEnv.serializerManager).thenReturn(serializerManager)
+      SparkEnv.set(mockEnv)
+      val mem = new TestMemoryManager(conf)
       mem.limit(memBytes)
       val taskM = new TaskMemoryManager(mem, 0)
       new UnsafeRowBuffer(taskM, Utils.createTempDir(), 1)
