diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/ValueState.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/ValueState.java index 92cfd0bd51df..560098fcf68d 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/state/ValueState.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/state/ValueState.java @@ -37,6 +37,15 @@ public interface ValueState extends ReadableState<@Nullable T>, State { @Nullable T read(); + /** + * Returns the current value of the state, or {@code defaultValue} if the value has never been + * written. + */ + default T read(T defaultValue) { + T value = read(); + return value == null ? defaultValue : value; + } + @Override ValueState readLater(); } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/state/ValueStateTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/state/ValueStateTest.java new file mode 100644 index 000000000000..3ec1ceb441b2 --- /dev/null +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/state/ValueStateTest.java @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.state; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; + +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ValueState}. */ +@RunWith(JUnit4.class) +public class ValueStateTest { + + @Test + public void testReadReturnsDefaultValueWhenStateIsEmpty() { + ValueState state = new TestValueState<>(); + assertEquals(Integer.valueOf(5), state.read(5)); + } + + @Test + public void testReadReturnsStoredValueWhenStateIsPresent() { + TestValueState state = new TestValueState<>(); + state.write(10); + assertEquals(Integer.valueOf(10), state.read(5)); + } + + @Test + public void testReadLaterReturnsSameState() { + ValueState state = new TestValueState<>(); + assertSame(state, state.readLater()); + } + + private static class TestValueState implements ValueState { + + private @Nullable T value; + + @Override + public void write(T input) { + value = input; + } + + @Override + public @Nullable T read() { + return value; + } + + @Override + public ValueState readLater() { + return this; + } + + @Override + public void clear() { + value = null; + } + } +}