@Override public void setup( StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) { this.container = containingTask; this.config = config; this.output = output; this.runtimeContext = new StreamingRuntimeContext( this, container.getEnvironment(), container.getAccumulatorMap()); stateKeySelector1 = config.getStatePartitioner(0, getUserCodeClassloader()); stateKeySelector2 = config.getStatePartitioner(1, getUserCodeClassloader()); try { TypeSerializer<Object> keySerializer = config.getStateKeySerializer(getUserCodeClassloader()); // if the keySerializer is null we still need to create the state backend // for the non-partitioned state features it provides, such as the state output streams String operatorIdentifier = getClass().getSimpleName() + "_" + config.getVertexID() + "_" + runtimeContext.getIndexOfThisSubtask(); stateBackend = container.createStateBackend(operatorIdentifier, keySerializer); } catch (Exception e) { throw new RuntimeException("Could not initialize state backend. ", e); } }
private static StreamTask<?, ?> createMockTaskWithTimer( final ScheduledExecutorService timerService, final Object lock) { StreamTask<?, ?> task = mock(StreamTask.class); when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>()); when(task.getName()).thenReturn("Test task name"); when(task.getExecutionConfig()).thenReturn(new ExecutionConfig()); when(task.getEnvironment()) .thenReturn( new MockEnvironment("MockTask", 3 * 1024 * 1024, new MockInputSplitProvider(), 1024)); when(task.getCheckpointLock()).thenReturn(lock); doAnswer( new Answer<Void>() { @Override public Void answer(InvocationOnMock invocationOnMock) throws Throwable { final Long timestamp = (Long) invocationOnMock.getArguments()[0]; final Triggerable target = (Triggerable) invocationOnMock.getArguments()[1]; timerService.schedule( new Callable<Object>() { @Override public Object call() throws Exception { synchronized (lock) { target.trigger(timestamp); } return null; } }, timestamp - System.currentTimeMillis(), TimeUnit.MILLISECONDS); return null; } }) .when(task) .registerTimer(anyLong(), any(Triggerable.class)); // ugly Java generic hacks to get the generic state backend into the mock @SuppressWarnings("unchecked") OngoingStubbing<StateBackend<?>> stubbing = (OngoingStubbing<StateBackend<?>>) (OngoingStubbing<?>) when(task.getStateBackend()); stubbing.thenReturn(MemoryStateBackend.defaultInstance()); return task; }
private static StreamTask<?, ?> createMockTask() { Configuration configuration = new Configuration(); configuration.setString(ConfigConstants.STATE_BACKEND, "jobmanager"); StreamTask<?, ?> task = mock(StreamTask.class); when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>()); when(task.getName()).thenReturn("Test task name"); when(task.getExecutionConfig()).thenReturn(new ExecutionConfig()); final TaskManagerRuntimeInfo mockTaskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class); when(mockTaskManagerRuntimeInfo.getConfiguration()).thenReturn(configuration); final Environment env = mock(Environment.class); when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 1, 0, 1, 0)); when(env.getUserClassLoader()) .thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader()); when(env.getMetricGroup()).thenReturn(new UnregisteredTaskMetricsGroup()); when(task.getEnvironment()).thenReturn(env); return task; }
private static StreamTask<?, ?> createMockTask() { StreamTask<?, ?> task = mock(StreamTask.class); when(task.getAccumulatorMap()).thenReturn(new HashMap<String, Accumulator<?, ?>>()); when(task.getName()).thenReturn("Test task name"); when(task.getExecutionConfig()).thenReturn(new ExecutionConfig()); Environment env = mock(Environment.class); when(env.getTaskInfo()).thenReturn(new TaskInfo("Test task name", 0, 1, 0)); when(env.getUserClassLoader()) .thenReturn(AggregatingAlignedProcessingTimeWindowOperatorTest.class.getClassLoader()); when(task.getEnvironment()).thenReturn(env); // ugly java generic hacks to get the state backend into the mock @SuppressWarnings("unchecked") OngoingStubbing<StateBackend<?>> stubbing = (OngoingStubbing<StateBackend<?>>) (OngoingStubbing<?>) when(task.getStateBackend()); stubbing.thenReturn(MemoryStateBackend.defaultInstance()); return task; }