Examples with StateNamespaceForTest org.apache.beam.runners.core.StateNamespaceForTest used on opensource projects

Example 6 with StateNamespaceForTest

use of org.apache.beam.runners.core.StateNamespaceForTest in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testBagStateWithUnderlying.

@Test
public void testBagStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<BagState<Integer>> valueTag = StateTags.bag("foo", VarIntCoder.of());
    BagState<Integer> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), emptyIterable());
    underlyingValue.add(1);
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    BagState<Integer> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(1));
    copyOnAccessState.add(4);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(4, 1));
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    BagState<Integer> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}

Example 7 with StateNamespaceForTest

use of org.apache.beam.runners.core.StateNamespaceForTest in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testGetWithEmpty.

@Test
public void testGetWithEmpty() {
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of());
    BagState<String> stringBag = internals.state(namespace, bagTag);
    assertThat(stringBag.read(), emptyIterable());
    stringBag.add("bar");
    stringBag.add("baz");
    assertThat(stringBag.read(), containsInAnyOrder("baz", "bar"));
    BagState<String> reReadStringBag = internals.state(namespace, bagTag);
    assertThat(reReadStringBag.read(), containsInAnyOrder("baz", "bar"));
}

Example 8 with StateNamespaceForTest

use of org.apache.beam.runners.core.StateNamespaceForTest in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testWatermarkHoldStateWithUnderlying.

@Test
public void testWatermarkHoldStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    TimestampCombiner timestampCombiner = TimestampCombiner.EARLIEST;
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<WatermarkHoldState> stateTag = StateTags.watermarkStateInternal("wmstate", timestampCombiner);
    WatermarkHoldState underlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), nullValue());
    underlyingValue.add(new Instant(250L));
    assertThat(underlyingValue.read(), equalTo(new Instant(250L)));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    WatermarkHoldState copyOnAccessState = internals.state(namespace, stateTag);
    assertThat(copyOnAccessState.read(), equalTo(new Instant(250L)));
    copyOnAccessState.add(new Instant(100L));
    assertThat(copyOnAccessState.read(), equalTo(new Instant(100L)));
    assertThat(underlyingValue.read(), equalTo(new Instant(250L)));
    copyOnAccessState.add(new Instant(500L));
    assertThat(copyOnAccessState.read(), equalTo(new Instant(100L)));
    WatermarkHoldState reReadUnderlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}

Example 9 with StateNamespaceForTest

use of org.apache.beam.runners.core.StateNamespaceForTest in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testSetStateWithUnderlying.

@Test
public void testSetStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<SetState<Integer>> valueTag = StateTags.set("foo", VarIntCoder.of());
    SetState<Integer> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), emptyIterable());
    underlyingValue.add(1);
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    SetState<Integer> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(1));
    copyOnAccessState.add(4);
    assertThat(copyOnAccessState.read(), containsInAnyOrder(4, 1));
    assertThat(underlyingValue.read(), containsInAnyOrder(1));
    SetState<Integer> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}

Example 10 with StateNamespaceForTest

use of org.apache.beam.runners.core.StateNamespaceForTest in project beam by apache.

the class CopyOnAccessInMemoryStateInternalsTest method testAccumulatorCombiningStateWithUnderlying.

@Test
public void testAccumulatorCombiningStateWithUnderlying() throws CannotProvideCoderException {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CombineFn<Long, long[], Long> sumLongFn = Sum.ofLongs();
    StateNamespace namespace = new StateNamespaceForTest("foo");
    CoderRegistry reg = pipeline.getCoderRegistry();
    StateTag<CombiningState<Long, long[], Long>> stateTag = StateTags.combiningValue("summer", sumLongFn.getAccumulatorCoder(reg, reg.getCoder(Long.class)), sumLongFn);
    GroupingState<Long, Long> underlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(0L));
    underlyingValue.add(1L);
    assertThat(underlyingValue.read(), equalTo(1L));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    GroupingState<Long, Long> copyOnAccessState = internals.state(namespace, stateTag);
    assertThat(copyOnAccessState.read(), equalTo(1L));
    copyOnAccessState.add(4L);
    assertThat(copyOnAccessState.read(), equalTo(5L));
    assertThat(underlyingValue.read(), equalTo(1L));
    GroupingState<Long, Long> reReadUnderlyingValue = underlying.state(namespace, stateTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}