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

Example 11 with StateNamespaceForTest

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

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithOnlyClearedValuesIsEmpty.

@Test
public void testCommitWithOnlyClearedValuesIsEmpty() {
    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("foo");
    stringBag.clear();
    internals.commit();
    assertThat(internals.isEmpty(), is(true));
}

Example 12 with StateNamespaceForTest

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

the class CopyOnAccessInMemoryStateInternalsTest method testMapStateWithUnderlying.

@Test
public void testMapStateWithUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<MapState<String, Integer>> valueTag = StateTags.map("foo", StringUtf8Coder.of(), VarIntCoder.of());
    MapState<String, Integer> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.entries().read(), emptyIterable());
    underlyingValue.put("hello", 1);
    assertThat(underlyingValue.get("hello").read(), equalTo(1));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    MapState<String, Integer> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.get("hello").read(), equalTo(1));
    copyOnAccessState.put("world", 4);
    assertThat(copyOnAccessState.get("hello").read(), equalTo(1));
    assertThat(copyOnAccessState.get("world").read(), equalTo(4));
    assertThat(underlyingValue.get("hello").read(), equalTo(1));
    assertNull(underlyingValue.get("world").read());
    MapState<String, Integer> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.entries().read(), equalTo(reReadUnderlyingValue.entries().read()));
}

Example 13 with StateNamespaceForTest

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

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithClearedInUnderlying.

@Test
public void testCommitWithClearedInUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    CopyOnAccessInMemoryStateInternals<String> secondUnderlying = spy(CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, secondUnderlying);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of());
    BagState<String> stringBag = underlying.state(namespace, bagTag);
    assertThat(stringBag.read(), emptyIterable());
    stringBag.add("bar");
    stringBag.add("baz");
    stringBag.clear();
    // We should not read through the cleared bag
    secondUnderlying.commit();
    // Should not be visible
    stringBag.add("foo");
    internals.commit();
    BagState<String> internalsStringBag = internals.state(namespace, bagTag);
    assertThat(internalsStringBag.read(), emptyIterable());
    verify(secondUnderlying, never()).state(namespace, bagTag);
    assertThat(internals.isEmpty(), is(false));
}

Example 14 with StateNamespaceForTest

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

the class CopyOnAccessInMemoryStateInternalsTest method testCommitWithAddedUnderlying.

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

Example 15 with StateNamespaceForTest

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

the class CopyOnAccessInMemoryStateInternalsTest method testGetWithPresentInUnderlying.

/**
   * Tests that retrieving state with an underlying StateInternals with an existing value returns
   * a value that initially has equal value to the provided state but can be modified without
   * modifying the existing state.
   */
@Test
public void testGetWithPresentInUnderlying() {
    CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null);
    StateNamespace namespace = new StateNamespaceForTest("foo");
    StateTag<ValueState<String>> valueTag = StateTags.value("foo", StringUtf8Coder.of());
    ValueState<String> underlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), nullValue(String.class));
    underlyingValue.write("bar");
    assertThat(underlyingValue.read(), equalTo("bar"));
    CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying);
    ValueState<String> copyOnAccessState = internals.state(namespace, valueTag);
    assertThat(copyOnAccessState.read(), equalTo("bar"));
    copyOnAccessState.write("baz");
    assertThat(copyOnAccessState.read(), equalTo("baz"));
    assertThat(underlyingValue.read(), equalTo("bar"));
    ValueState<String> reReadUnderlyingValue = underlying.state(namespace, valueTag);
    assertThat(underlyingValue.read(), equalTo(reReadUnderlyingValue.read()));
}