@Override public ListState<String> answer(InvocationOnMock invocationOnMock) throws Throwable { ListStateDescriptor<String> descr = (ListStateDescriptor<String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(ListStateDescriptor.class));
@Override public MapState<Integer, String> answer(InvocationOnMock invocationOnMock) throws Throwable { MapStateDescriptor<Integer, String> descr = (MapStateDescriptor<Integer, String>) invocationOnMock.getArguments()[2]; AbstractKeyedStateBackend<Integer> backend = new MemoryStateBackend().createKeyedStateBackend( new DummyEnvironment("test_task", 1, 0), new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); backend.setCurrentKey(0); return backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, descr); } }).when(keyedStateBackend).getPartitionedState(Matchers.any(), any(TypeSerializer.class), any(MapStateDescriptor.class));
public TriggerTestHarness( Trigger<T, W> trigger, TypeSerializer<W> windowSerializer) throws Exception { this.trigger = trigger; this.windowSerializer = windowSerializer; // we only ever use one key, other tests make sure that windows work across different // keys DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0); MemoryStateBackend backend = new MemoryStateBackend(); @SuppressWarnings("unchecked") HeapKeyedStateBackend<Integer> stateBackend = (HeapKeyedStateBackend<Integer>) backend.createKeyedStateBackend(dummyEnv, new JobID(), "test_op", IntSerializer.INSTANCE, 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); this.stateBackend = stateBackend; this.stateBackend.setCurrentKey(KEY); this.internalTimerService = new TestInternalTimerService<>(new KeyContext() { @Override public void setCurrentKey(Object key) { // ignore } @Override public Object getCurrentKey() { return KEY; } }); }
private static KeyedStateBackend<ByteBuffer> getKeyedStateBackend(int numberOfKeyGroups, KeyGroupRange keyGroupRange) { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), numberOfKeyGroups, keyGroupRange, new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey(ByteBuffer.wrap( CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1"))); return keyedStateBackend; } catch (Exception e) { throw new RuntimeException(e); } }
private static KeyedStateBackend<ByteBuffer> getKeyedStateBackend( int numberOfKeyGroups, KeyGroupRange keyGroupRange) { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), numberOfKeyGroups, keyGroupRange, new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey( ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1"))); return keyedStateBackend; } catch (Exception e) { throw new RuntimeException(e); } } }
private static KeyedStateBackend<ByteBuffer> getKeyedStateBackend( int numberOfKeyGroups, KeyGroupRange keyGroupRange) { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), numberOfKeyGroups, keyGroupRange, new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey( ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "1"))); return keyedStateBackend; } catch (Exception e) { throw new RuntimeException(e); } } }
@Override protected StateInternals createStateInternals() { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey( ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Hello"))); return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of()); } catch (Exception e) { throw new RuntimeException(e); } } }
@Override protected StateInternals createStateInternals() { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey( ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Hello"))); return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of()); } catch (Exception e) { throw new RuntimeException(e); } }
@Override protected StateInternals createStateInternals() { MemoryStateBackend backend = new MemoryStateBackend(); try { AbstractKeyedStateBackend<ByteBuffer> keyedStateBackend = backend.createKeyedStateBackend( new DummyEnvironment("test", 1, 0), new JobID(), "test_op", new GenericTypeInfo<>(ByteBuffer.class).createSerializer(new ExecutionConfig()), 1, new KeyGroupRange(0, 0), new KvStateRegistry().createTaskRegistry(new JobID(), new JobVertexID())); keyedStateBackend.setCurrentKey( ByteBuffer.wrap(CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Hello"))); return new FlinkStateInternals<>(keyedStateBackend, StringUtf8Coder.of()); } catch (Exception e) { throw new RuntimeException(e); } } }