Consecutive GetState with empty (#335)

restatedev · May 31, 2024 · 01dab68 · 01dab68
1 parent 957c7b6
commit 01dab68
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 2 deletions.
diff --git a/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/EagerStateTest.kt b/sdk-api-kotlin/src/test/kotlin/dev/restate/sdk/kotlin/EagerStateTest.kt
@@ -57,4 +57,10 @@ class EagerStateTest : EagerStateTestSuite() {
       testDefinitionForVirtualObject("ListKeys") { ctx, _: Unit ->
         ctx.stateKeys().joinToString(separator = ",")
       }
+
+  override fun consecutiveGetWithEmpty(): TestInvocationBuilder =
+      testDefinitionForVirtualObject("ConsecutiveGetWithEmpty") { ctx, _: Unit ->
+        assertThat(ctx.get(StateKey.of("key-0", TestSerdes.STRING))).isNull()
+        assertThat(ctx.get(StateKey.of("key-0", TestSerdes.STRING))).isNull()
+      }
 }
diff --git a/sdk-api/src/test/java/dev/restate/sdk/EagerStateTest.java b/sdk-api/src/test/java/dev/restate/sdk/EagerStateTest.java
@@ -85,4 +85,17 @@ protected TestInvocationBuilder listKeys() {
         JsonSerdes.STRING,
         (ctx, input) -> String.join(",", ctx.stateKeys()));
   }
+
+  @Override
+  protected TestInvocationBuilder consecutiveGetWithEmpty() {
+    return testDefinitionForVirtualObject(
+        "ConsecutiveGetWithEmpty",
+        Serde.VOID,
+        Serde.VOID,
+        (ctx, input) -> {
+          assertThat(ctx.get(StateKey.of("key-0", JsonSerdes.STRING))).isEmpty();
+          assertThat(ctx.get(StateKey.of("key-0", JsonSerdes.STRING))).isEmpty();
+          return null;
+        });
+  }
 }
diff --git a/sdk-core/src/main/java/dev/restate/sdk/core/Entries.java b/sdk-core/src/main/java/dev/restate/sdk/core/Entries.java
@@ -140,7 +140,11 @@ public Result<ByteBuffer> parseCompletionResult(CompletionMessage actual) {
     @Override
     void updateUserStateStoreWithEntry(
         GetStateEntryMessage expected, UserStateStore userStateStore) {
-      userStateStore.set(expected.getKey(), expected.getValue().asReadOnlyByteBuffer());
+      if (expected.hasEmpty()) {
+        userStateStore.clear(expected.getKey());
+      } else {
+        userStateStore.set(expected.getKey(), expected.getValue().asReadOnlyByteBuffer());
+      }
     }
 
     @Override

diff --git a/sdk-core/src/test/java/dev/restate/sdk/core/EagerStateTestSuite.java b/sdk-core/src/test/java/dev/restate/sdk/core/EagerStateTestSuite.java
@@ -15,6 +15,7 @@
 import com.google.protobuf.MessageLite;
 import dev.restate.generated.service.protocol.Protocol;
 import dev.restate.generated.service.protocol.Protocol.ClearAllStateEntryMessage;
+import dev.restate.generated.service.protocol.Protocol.Empty;
 import java.util.Map;
 import java.util.stream.Stream;
 
@@ -32,6 +33,8 @@ public abstract class EagerStateTestSuite implements TestSuite {
 
   protected abstract TestInvocationBuilder listKeys();
 
+  protected abstract TestInvocationBuilder consecutiveGetWithEmpty();
+
   private static final Map.Entry<String, String> STATE_FRANCESCO = entry("STATE", "Francesco");
   private static final Map.Entry<String, String> ANOTHER_STATE_FRANCESCO =
       entry("ANOTHER_STATE", "Francesco");
@@ -165,6 +168,23 @@ public Stream<TestDefinition> definitions() {
                 INPUT_TILL,
                 Protocol.GetStateKeysEntryMessage.newBuilder().setValue(stateKeys("3", "2", "1")))
             .expectingOutput(outputMessage("3,2,1"), END_MESSAGE)
-            .named("With replayed list"));
+            .named("With replayed list"),
+        this.consecutiveGetWithEmpty()
+            .withInput(startMessage(1).setPartialState(false), inputMessage())
+            .expectingOutput(
+                getStateMessage("key-0").setEmpty(Empty.getDefaultInstance()),
+                getStateMessage("key-0").setEmpty(Empty.getDefaultInstance()),
+                outputMessage(),
+                END_MESSAGE),
+        this.consecutiveGetWithEmpty()
+            .withInput(
+                startMessage(2).setPartialState(false),
+                inputMessage(),
+                getStateMessage("key-0").setEmpty(Empty.getDefaultInstance()))
+            .expectingOutput(
+                getStateMessage("key-0").setEmpty(Empty.getDefaultInstance()),
+                outputMessage(),
+                END_MESSAGE)
+            .named("With replay of the first get"));
   }
 }