From 01dab6877eaeec865322fc8d4d4473573f2b5d32 Mon Sep 17 00:00:00 2001
From: Francesco Guardiani <me@slinkydeveloper.com>
Date: Fri, 31 May 2024 16:52:14 +0200
Subject: [PATCH] Consecutive GetState with empty (#335)

---
 .../dev/restate/sdk/kotlin/EagerStateTest.kt  |  6 +++++
 .../java/dev/restate/sdk/EagerStateTest.java  | 13 +++++++++++
 .../java/dev/restate/sdk/core/Entries.java    |  6 ++++-
 .../restate/sdk/core/EagerStateTestSuite.java | 22 ++++++++++++++++++-
 4 files changed, 45 insertions(+), 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
index 897408f6..5e9fb2f2 100644
--- 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
index f53b1b88..4a0dd695 100644
--- 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
index c7502460..b6f04dfd 100644
--- 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
index 285d6319..2925db24 100644
--- 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"));
   }
 }