1 files changed, 66 insertions, 0 deletions
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
index f151c5aa54..776cf2da6c 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skyframe.serialization;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.common.truth.Truth.assertWithMessage;
 import static com.google.devtools.build.lib.testutil.MoreAsserts.assertThrows;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.when;
@@ -23,9 +24,11 @@ import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec.MemoizationStrategy;
+import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.testutils.TestUtils;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -128,6 +131,69 @@ public class SerializationContextTest {
   }
 
   @Test
+  public void memoizingStringPrefixes_EndtoEnd() throws Exception {
+    ObjectCodec<String> codec = StringCodecs.simple();
+    Memoizer.Serializer serializer = new Memoizer.Serializer();
+    SerializationContext serializationContext = Mockito.mock(SerializationContext.class);
+    ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+    CodedOutputStream codedOutputStream = CodedOutputStream.newInstance(byteArrayOutputStream);
+
+    // Serialize all the things.
+    serializer.serialize(serializationContext, "string", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "string2", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "string2", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "strip", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "banana", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "", codec, codedOutputStream);
+    serializer.serialize(serializationContext, "peach", codec, codedOutputStream);
+
+    // SerializationContext not used for simple string serialization.
+    Mockito.verifyZeroInteractions(serializationContext);
+
+    // Flush outputs and assert not too much data was written.
+    codedOutputStream.flush();
+    byteArrayOutputStream.flush();
+    byte[] bytes = byteArrayOutputStream.toByteArray();
+    int stringOverhead = 1;
+    int prefixOverhead = 1;
+    // Every string but the first has prefixOverhead, every string has stringOverhead.
+    assertThat(bytes.length)
+        .isEqualTo(
+            "string".length()
+                + "2".length()
+                + "p".length()
+                + "banana".length()
+                + "peach".length()
+                + 7 * stringOverhead
+                + 6 * prefixOverhead);
+
+    // Prepare inputs.
+    CodedInputStream codedInputStream = CodedInputStream.newInstance(bytes);
+    Memoizer.Deserializer deserializer = new Memoizer.Deserializer();
+    DeserializationContext deserializationContext = Mockito.mock(DeserializationContext.class);
+
+    // Deserialize and assert fidelity.
+    assertThat(deserializer.deserialize(deserializationContext, codec, codedInputStream))
+        .isEqualTo("string");
+    String returnedString =
+        deserializer.deserialize(deserializationContext, codec, codedInputStream);
+    assertThat(returnedString).isEqualTo("string2");
+    String newReturnedString =
+        deserializer.deserialize(deserializationContext, codec, codedInputStream);
+    assertThat(newReturnedString).isEqualTo("string2");
+    assertWithMessage("Same string twice in a row should be the same object")
+        .that(newReturnedString)
+        .isSameAs(returnedString);
+    assertThat(deserializer.deserialize(deserializationContext, codec, codedInputStream))
+        .isEqualTo("strip");
+    assertThat(deserializer.deserialize(deserializationContext, codec, codedInputStream))
+        .isEqualTo("banana");
+    assertThat(deserializer.deserialize(deserializationContext, codec, codedInputStream)).isEmpty();
+    assertThat(deserializer.deserialize(deserializationContext, codec, codedInputStream))
+        .isEqualTo("peach");
+  }
+
+  @Test
   public void startMemoizingIsIdempotent() throws IOException, SerializationException {
     ObjectCodecRegistry registry =
         ObjectCodecRegistry.newBuilder()