TaskFactoryUtil { public static TaskFactory finalizeTaskFactory(TaskFactory factory, ExecutorService taskThreadPool) { validateFactory(factory); boolean isAsyncTaskClass = factory instanceof AsyncStreamTaskFactory; if (isAsyncTaskClass) { log.info("Got an AsyncStreamTask implementation."); return factory; } boolean isStreamOperatorTaskClass = factory instanceof StreamOperatorTaskFactory; if (isStreamOperatorTaskClass) { log.info("Adapting StreamOperatorTaskFactory to inject container thread pool"); return (AsyncStreamTaskFactory) () -> { StreamOperatorTask operatorTask = (StreamOperatorTask) factory.createInstance(); operatorTask.setTaskThreadPool(taskThreadPool); return operatorTask; }; } log.info("Converting StreamTask to AsyncStreamTaskAdapter"); return (AsyncStreamTaskFactory) () -> new AsyncStreamTaskAdapter(((StreamTaskFactory) factory).createInstance(), taskThreadPool); }  static TaskFactory getTaskFactory(ApplicationDescriptorImpl<? extends ApplicationDescriptor> appDesc); static TaskFactory getTaskFactory(String taskClassName); static TaskFactory finalizeTaskFactory(TaskFactory factory, ExecutorService taskThreadPool);  }
@Test public void testFinalizeTaskFactory() throws NoSuchFieldException, IllegalAccessException { TaskFactory mockFactory = mock(TaskFactory.class); try { TaskFactoryUtil.finalizeTaskFactory(mockFactory, null); fail("Should have failed with validation"); } catch (SamzaException se) { } StreamTaskFactory mockStreamFactory = mock(StreamTaskFactory.class); ExecutorService mockThreadPool = mock(ExecutorService.class); TaskFactory retFactory = TaskFactoryUtil.finalizeTaskFactory(mockStreamFactory, mockThreadPool); assertTrue(retFactory instanceof AsyncStreamTaskFactory); assertTrue(((AsyncStreamTaskFactory) retFactory).createInstance() instanceof AsyncStreamTaskAdapter); AsyncStreamTaskAdapter taskAdapter = (AsyncStreamTaskAdapter) ((AsyncStreamTaskFactory) retFactory).createInstance(); Field executorSrvFld = AsyncStreamTaskAdapter.class.getDeclaredField("executor"); executorSrvFld.setAccessible(true); ExecutorService executor = (ExecutorService) executorSrvFld.get(taskAdapter); assertEquals(executor, mockThreadPool); AsyncStreamTaskFactory mockAsyncStreamFactory = mock(AsyncStreamTaskFactory.class); retFactory = TaskFactoryUtil.finalizeTaskFactory(mockAsyncStreamFactory, null); assertEquals(retFactory, mockAsyncStreamFactory); }
@Test public void testFinalizeTaskFactoryForStreamOperatorTask() { TaskFactory mockFactory = mock(StreamOperatorTaskFactory.class); StreamOperatorTask mockStreamOperatorTask = mock(StreamOperatorTask.class); when(mockFactory.createInstance()) .thenReturn(mockStreamOperatorTask); ExecutorService mockThreadPool = mock(ExecutorService.class); TaskFactory finalizedFactory = TaskFactoryUtil.finalizeTaskFactory(mockFactory, mockThreadPool); finalizedFactory.createInstance(); verify(mockStreamOperatorTask, times(1)).setTaskThreadPool(eq(mockThreadPool)); }