/** * Test that forks work correctly when the operator picks all outgoing forks */ @Test(dataProvider = "stateOverrides") public void testForkCorrectnessIdentity(State overrides) throws Exception { // Create a TaskState TaskState taskState = getEmptyTestTaskState("testForkTaskId"); taskState.addAll(overrides); int numRecords = 100; int numForks = 5; // Identity Fork Operator looks for number of forks in work unit state. taskState.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, "" + numForks); ForkOperator mockForkOperator = new IdentityForkOperator(); ArrayList<ArrayList<Object>> recordCollectors = runTaskAndGetResults(taskState, numRecords, numForks, mockForkOperator); // Check that we got the right records in the collectors int recordsPerFork = numRecords; for (int forkNumber=0; forkNumber < numForks; ++ forkNumber) { ArrayList<Object> forkRecords = recordCollectors.get(forkNumber); Assert.assertEquals(forkRecords.size(), recordsPerFork); for (int j=0; j < recordsPerFork; ++j) { Object forkRecord = forkRecords.get(j); Assert.assertEquals((String) forkRecord, "" + j); } } }
private TaskContext getMockTaskContext(ArrayList<Object> recordCollector, Extractor mockExtractor) throws Exception { TaskState taskState = getStreamingTaskState(); // Create a mock RowLevelPolicyChecker RowLevelPolicyChecker mockRowLevelPolicyChecker = new RowLevelPolicyChecker(Lists.newArrayList(), "stateId", FileSystem.getLocal(new Configuration())); WatermarkStorage mockWatermarkStorage = new MockWatermarkStorage(); // Create a mock TaskPublisher TaskPublisher mockTaskPublisher = mock(TaskPublisher.class); when(mockTaskPublisher.canPublish()).thenReturn(TaskPublisher.PublisherState.SUCCESS); // Create a mock TaskContext TaskContext mockTaskContext = mock(TaskContext.class); when(mockTaskContext.getTaskMetrics()).thenReturn(TaskMetrics.get(taskState)); when(mockTaskContext.getExtractor()).thenReturn(mockExtractor); when(mockTaskContext.getRawSourceExtractor()).thenReturn(mockExtractor); when(mockTaskContext.getWatermarkStorage()).thenReturn(mockWatermarkStorage); when(mockTaskContext.getForkOperator()).thenReturn(new IdentityForkOperator()); when(mockTaskContext.getTaskState()).thenReturn(taskState); when(mockTaskContext.getTaskPublisher(any(TaskState.class), any(TaskLevelPolicyCheckResults.class))) .thenReturn(mockTaskPublisher); when(mockTaskContext.getRowLevelPolicyChecker()).thenReturn(mockRowLevelPolicyChecker); when(mockTaskContext.getRowLevelPolicyChecker(anyInt())).thenReturn(mockRowLevelPolicyChecker); when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt())).thenReturn(mock(TaskLevelPolicyChecker.class)); when(mockTaskContext.getDataWriterBuilder(anyInt(), anyInt())).thenReturn(new TestStreamingDataWriterBuilder(recordCollector)); return mockTaskContext; }
@Test public void testForkMethods() { ForkOperator<String, String> dummyForkOperator = new IdentityForkOperator<String, String>(); WorkUnitState workUnitState = new WorkUnitState(); workUnitState.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, 2); List<Boolean> schemas = dummyForkOperator.forkSchema(workUnitState, ""); Assert.assertEquals(schemas, Arrays.asList(true, true)); List<Boolean> records = dummyForkOperator.forkDataRecord(workUnitState, ""); Assert.assertEquals(records, Arrays.asList(true, true)); Assert.assertEquals(dummyForkOperator.getBranches(workUnitState), 2); workUnitState.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, 0); schemas = dummyForkOperator.forkSchema(workUnitState, ""); Assert.assertTrue(schemas.isEmpty()); records = dummyForkOperator.forkDataRecord(workUnitState, ""); Assert.assertTrue(records.isEmpty()); Assert.assertEquals(dummyForkOperator.getBranches(workUnitState), 0); } }
when(mockTaskContext.getForkOperator()).thenReturn(new IdentityForkOperator()); when(mockTaskContext.getTaskState()).thenReturn(taskState); when(mockTaskContext.getConverters()).thenReturn(converters);
when(mockTaskContext.getTaskMetrics()).thenReturn(TaskMetrics.get(taskState)); when(mockTaskContext.getExtractor()).thenReturn(new FailOnceExtractor()); when(mockTaskContext.getForkOperator()).thenReturn(new IdentityForkOperator()); when(mockTaskContext.getTaskState()).thenReturn(taskState); when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt()))
when(mockTaskContext.getForkOperator()).thenReturn(new IdentityForkOperator()); when(mockTaskContext.getTaskState()).thenReturn(taskState); when(mockTaskContext.getConverters()).thenReturn(Lists.newArrayList(converter));