/** * Get a pre-fork {@link RowLevelPolicyChecker} for executing row-level * {@link org.apache.gobblin.qualitychecker.row.RowLevelPolicy}. * * @return a {@link RowLevelPolicyChecker} */ public RowLevelPolicyChecker getRowLevelPolicyChecker() throws Exception { return getRowLevelPolicyChecker(-1); }
this.rowChecker = closer.register(this.taskContext.getRowLevelPolicyChecker()); } catch (Exception e) { try {
private TaskContext getMockTaskContext(TaskState taskState, Extractor mockExtractor, ArrayList<ArrayList<Object>> writerCollectors, ForkOperator mockForkOperator) throws Exception { int numForks = writerCollectors.size(); // Create a mock RowLevelPolicyChecker RowLevelPolicyChecker mockRowLevelPolicyChecker = spy(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockRowLevelPolicyChecker.executePolicies(any(Object.class), any(RowLevelPolicyCheckResults.class))) .thenReturn(true); when(mockRowLevelPolicyChecker.getFinalState()).thenReturn(new State()); // 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.getExtractor()).thenReturn(mockExtractor); when(mockTaskContext.getRawSourceExtractor()).thenReturn(mockExtractor); when(mockTaskContext.getForkOperator()).thenReturn(mockForkOperator); 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)); for (int i =0; i < numForks; ++i) { when(mockTaskContext.getDataWriterBuilder(numForks, i)).thenReturn(new RecordCollectingWriterBuilder(writerCollectors.get(i))); } return mockTaskContext; }
when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt())) .thenReturn(mock(TaskLevelPolicyChecker.class)); when(mockTaskContext.getRowLevelPolicyChecker()). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockTaskContext.getRowLevelPolicyChecker(anyInt())). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockTaskContext.getDataWriterBuilder(anyInt(), anyInt())).thenReturn(writer);
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; }
when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt())) .thenReturn(mock(TaskLevelPolicyChecker.class)); when(mockTaskContext.getRowLevelPolicyChecker()). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockTaskContext.getRowLevelPolicyChecker(anyInt())). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockTaskContext.getDataWriterBuilder(anyInt(), anyInt())).thenReturn(writer);
when(mockTaskContext.getTaskLevelPolicyChecker(any(TaskState.class), anyInt())) .thenReturn(mock(TaskLevelPolicyChecker.class)); when(mockTaskContext.getRowLevelPolicyChecker()). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration()))); when(mockTaskContext.getRowLevelPolicyChecker(anyInt())). thenReturn(new RowLevelPolicyChecker(Lists.newArrayList(), "ss", FileSystem.getLocal(new Configuration())));
this.closer.register(new MultiConverter(this.taskContext.getConverters(this.index, this.forkTaskState))); this.convertedSchema = Optional.fromNullable(this.converter.convertSchema(schema, this.taskState)); this.rowLevelPolicyChecker = this.closer.register(this.taskContext.getRowLevelPolicyChecker(this.index)); this.rowLevelPolicyCheckingResult = new RowLevelPolicyCheckResults();
/** * Get a pre-fork {@link RowLevelPolicyChecker} for executing row-level * {@link org.apache.gobblin.qualitychecker.row.RowLevelPolicy}. * * @return a {@link RowLevelPolicyChecker} */ public RowLevelPolicyChecker getRowLevelPolicyChecker() throws Exception { return getRowLevelPolicyChecker(-1); }
this.rowChecker = closer.register(this.taskContext.getRowLevelPolicyChecker()); } catch (Exception e) { try {
this.closer.register(new MultiConverter(this.taskContext.getConverters(this.index, this.forkTaskState))); this.convertedSchema = Optional.fromNullable(this.converter.convertSchema(schema, this.taskState)); this.rowLevelPolicyChecker = this.closer.register(this.taskContext.getRowLevelPolicyChecker(this.index)); this.rowLevelPolicyCheckingResult = new RowLevelPolicyCheckResults();