@Test public void testHandleInterruptedException() throws Exception { HiveConf hiveConf = new HiveConf(); SparkTask sparkTask = new SparkTask(); sparkTask.setWork(mock(SparkWork.class)); DriverContext mockDriverContext = mock(DriverContext.class); QueryState mockQueryState = mock(QueryState.class); when(mockQueryState.getConf()).thenReturn(hiveConf); sparkTask.initialize(mockQueryState, null, mockDriverContext, null); SparkJobStatus mockSparkJobStatus = mock(SparkJobStatus.class); when(mockSparkJobStatus.getMonitorError()).thenReturn(new InterruptedException()); SparkSession mockSparkSession = mock(SparkSession.class); SparkJobRef mockSparkJobRef = mock(SparkJobRef.class); when(mockSparkJobRef.monitorJob()).thenReturn(2); when(mockSparkJobRef.getSparkJobStatus()).thenReturn(mockSparkJobStatus); when(mockSparkSession.submit(any(), any())).thenReturn(mockSparkJobRef); SessionState.start(hiveConf); SessionState.get().setSparkSession(mockSparkSession); sparkTask.execute(mockDriverContext); verify(mockSparkJobRef, atLeastOnce()).cancelJob(); when(mockSparkJobStatus.getMonitorError()).thenReturn( new HiveException(new InterruptedException())); sparkTask.execute(mockDriverContext); verify(mockSparkJobRef, atLeastOnce()).cancelJob(); }
@Test public void removeEmptySparkTask() { SparkTask grandpa = new SparkTask(); SparkWork grandpaWork = new SparkWork("grandpa"); grandpaWork.add(new MapWork()); grandpa.setWork(grandpaWork); SparkTask parent = new SparkTask(); SparkWork parentWork = new SparkWork("parent"); parentWork.add(new MapWork()); parent.setWork(parentWork); SparkTask child1 = new SparkTask(); SparkWork childWork1 = new SparkWork("child1"); childWork1.add(new MapWork()); child1.setWork(childWork1); grandpa.addDependentTask(parent); parent.addDependentTask(child1); Assert.assertEquals(grandpa.getChildTasks().size(), 1); Assert.assertEquals(child1.getParentTasks().size(), 1); if (isEmptySparkWork(parent.getWork())) { SparkUtilities.removeEmptySparkTask(parent); } Assert.assertEquals(grandpa.getChildTasks().size(), 0); Assert.assertEquals(child1.getParentTasks().size(), 0); }