SerializationTester uses an addDependency method instead of

(Des|S)erializationContext factories.

This is consistent as ObjectCodecs will eventually create context objects if needed.

PiperOrigin-RevId: 186900662

2 files changed, 24 insertions, 47 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
index 4d31a033be..db996325f2 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
@@ -36,7 +36,8 @@ public class ObjectCodecs {
    * Creates an instance using the supplied {@link ObjectCodecRegistry} for looking up {@link
    * ObjectCodec}s.
    */
-  ObjectCodecs(ObjectCodecRegistry codecRegistry, ImmutableMap<Class<?>, Object> dependencies) {
+  public ObjectCodecs(
+      ObjectCodecRegistry codecRegistry, ImmutableMap<Class<?>, Object> dependencies) {
     this.codecRegistry = codecRegistry;
     serializationContext = new SerializationContext(codecRegistry, dependencies);
     deserializationContext = new DeserializationContext(codecRegistry, dependencies);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
index d89a428603..ee45a822c1 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializationTester.java
@@ -21,15 +21,11 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.devtools.build.lib.skyframe.serialization.DeserializationContext;
-import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
+import com.google.devtools.build.lib.skyframe.serialization.AutoRegistry;
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
-import com.google.protobuf.CodedInputStream;
-import com.google.protobuf.CodedOutputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.IOException;
+import com.google.protobuf.ByteString;
 import java.util.Random;
-import java.util.function.Supplier;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -57,10 +53,7 @@ public class SerializationTester {
   }
 
   private final ImmutableList<Object> subjects;
-  private Supplier<SerializationContext> writeContextFactory =
-      () -> new SerializationContext(ImmutableMap.of());
-  private Supplier<DeserializationContext> readContextFactory =
-      () -> new DeserializationContext(ImmutableMap.of());
+  private final ImmutableMap.Builder<Class<?>, Object> dependenciesBuilder;
 
   @SuppressWarnings("rawtypes")
   private VerificationFunction verificationFunction =
@@ -71,17 +64,11 @@ public class SerializationTester {
   public SerializationTester(Object... subjects) {
     Preconditions.checkArgument(subjects.length > 0);
     this.subjects = ImmutableList.copyOf(subjects);
+    this.dependenciesBuilder = ImmutableMap.builder();
   }
 
-  public SerializationTester setWriteContextFactory(
-      Supplier<SerializationContext> writeContextFactory) {
-    this.writeContextFactory = writeContextFactory;
-    return this;
-  }
-
-  public SerializationTester setReadContextFactory(
-      Supplier<DeserializationContext> readContextFactory) {
-    this.readContextFactory = readContextFactory;
+  public <D> SerializationTester addDependency(Class<? super D> type, D dependency) {
+    dependenciesBuilder.put(type, dependency);
     return this;
   }
 
@@ -98,21 +85,22 @@ public class SerializationTester {
   }
 
   public void runTests() throws Exception {
-    testSerializeDeserialize();
-    testStableSerialization();
-    testDeserializeJunkData();
+    ObjectCodecs codecs = new ObjectCodecs(AutoRegistry.get(), dependenciesBuilder.build());
+    testSerializeDeserialize(codecs);
+    testStableSerialization(codecs);
+    testDeserializeJunkData(codecs);
   }
 
   /** Runs serialization/deserialization tests. */
   @SuppressWarnings("unchecked")
-  private void testSerializeDeserialize() throws Exception {
+  private void testSerializeDeserialize(ObjectCodecs codecs) throws Exception {
     Stopwatch timer = Stopwatch.createStarted();
     int totalBytes = 0;
     for (int i = 0; i < repetitions; ++i) {
       for (Object subject : subjects) {
-        byte[] serialized = toBytes(subject);
-        totalBytes += serialized.length;
-        Object deserialized = fromBytes(serialized);
+        ByteString serialized = codecs.serialize(subject);
+        totalBytes += serialized.size();
+        Object deserialized = codecs.deserialize(serialized);
         verificationFunction.verifyDeserialized(subject, deserialized);
       }
     }
@@ -126,41 +114,29 @@ public class SerializationTester {
   }
 
   /** Runs serialized bytes stability tests. */
-  private void testStableSerialization() throws IOException, SerializationException {
+  private void testStableSerialization(ObjectCodecs codecs) throws SerializationException {
     for (Object subject : subjects) {
-      byte[] serialized = toBytes(subject);
-      Object deserialized = fromBytes(serialized);
-      byte[] reserialized = toBytes(deserialized);
+      ByteString serialized = codecs.serialize(subject);
+      Object deserialized = codecs.deserialize(serialized);
+      ByteString reserialized = codecs.serialize(deserialized);
       assertThat(reserialized).isEqualTo(serialized);
     }
   }
 
   /** Runs junk-data recognition tests. */
-  private void testDeserializeJunkData() {
+  private static void testDeserializeJunkData(ObjectCodecs codecs) {
     Random rng = new Random(0);
     for (int i = 0; i < DEFAULT_JUNK_INPUTS; ++i) {
       byte[] junkData = new byte[rng.nextInt(JUNK_LENGTH_UPPER_BOUND)];
       rng.nextBytes(junkData);
       try {
-        readContextFactory.get().deserialize(CodedInputStream.newInstance(junkData));
+        codecs.deserialize(ByteString.copyFrom(junkData));
         // OK. Junk string was coincidentally parsed.
-      } catch (IOException | SerializationException e) {
+      } catch (SerializationException e) {
         // OK. Deserialization of junk failed.
         return;
       }
     }
     assert_().fail("all junk was parsed successfully");
   }
-
-  private Object fromBytes(byte[] bytes) throws IOException, SerializationException {
-    return readContextFactory.get().deserialize(CodedInputStream.newInstance(bytes));
-  }
-
-  private byte[] toBytes(Object subject) throws IOException, SerializationException {
-    ByteArrayOutputStream bytes = new ByteArrayOutputStream();
-    CodedOutputStream codedOut = CodedOutputStream.newInstance(bytes);
-    writeContextFactory.get().serialize(subject, codedOut);
-    codedOut.flush();
-    return bytes.toByteArray();
-  }
 }