1 files changed, 22 insertions, 4 deletions
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
index 82d4c36baa..a0f5873ae6 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
@@ -20,8 +20,10 @@ import static org.junit.Assert.fail;
 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.ObjectCodec;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationContext;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
 import com.google.protobuf.CodedInputStream;
 import java.io.IOException;
@@ -55,17 +57,23 @@ public class ObjectCodecTester<T> {
 
   private final ObjectCodec<T> underTest;
   private final ImmutableList<T> subjects;
+  private final SerializationContext writeContext;
+  private final DeserializationContext readContext;
   private final boolean skipBadDataTest;
   private final VerificationFunction<T> verificationFunction;
 
   private ObjectCodecTester(
       ObjectCodec<T> underTest,
       ImmutableList<T> subjects,
+      SerializationContext writeContext,
+      DeserializationContext readContext,
       boolean skipBadDataTest,
       VerificationFunction<T> verificationFunction) {
     this.underTest = underTest;
     Preconditions.checkState(!subjects.isEmpty(), "No subjects provided");
     this.subjects = subjects;
+    this.writeContext = writeContext;
+    this.readContext = readContext;
     this.skipBadDataTest = skipBadDataTest;
     this.verificationFunction = verificationFunction;
   }
@@ -111,8 +119,7 @@ public class ObjectCodecTester<T> {
   void testDeserializeJunkData() {
     try {
       underTest.deserialize(
-          DeserializationContext.create(),
-          CodedInputStream.newInstance("junk".getBytes(StandardCharsets.UTF_8)));
+          readContext, CodedInputStream.newInstance("junk".getBytes(StandardCharsets.UTF_8)));
       fail("Expected exception");
     } catch (SerializationException | IOException e) {
       // Expected.
@@ -120,17 +127,19 @@ public class ObjectCodecTester<T> {
   }
 
   private T fromBytes(byte[] bytes) throws SerializationException, IOException {
-    return TestUtils.fromBytes(underTest, bytes);
+    return TestUtils.fromBytes(readContext, underTest, bytes);
   }
 
   private byte[] toBytes(T subject) throws IOException, SerializationException {
-    return TestUtils.toBytes(underTest, subject);
+    return TestUtils.toBytes(writeContext, underTest, subject);
   }
 
   /** Builder for {@link ObjectCodecTester}. */
   public static class Builder<T> {
     private final ObjectCodec<T> underTest;
     private final ImmutableList.Builder<T> subjectsBuilder = ImmutableList.builder();
+    private final ImmutableMap.Builder<Class<?>, Object> dependenciesBuilder =
+        ImmutableMap.builder();
     private boolean skipBadDataTest = false;
     private VerificationFunction<T> verificationFunction =
         (original, deserialized) -> assertThat(deserialized).isEqualTo(original);
@@ -151,6 +160,12 @@ public class ObjectCodecTester<T> {
       return this;
     }
 
+    /** Add subjects to be tested for serialization/deserialization. */
+    public final <D> Builder<T> addDependency(Class<? super D> type, D dependency) {
+      dependenciesBuilder.put(type, dependency);
+      return this;
+    }
+
     /**
      * Skip tests that check for the ability to detect bad data. This may be useful for simpler
      * codecs which don't do any error verification.
@@ -180,9 +195,12 @@ public class ObjectCodecTester<T> {
      * individually.
      */
     ObjectCodecTester<T> build() {
+      ImmutableMap<Class<?>, Object> dependencies = dependenciesBuilder.build();
       return new ObjectCodecTester<>(
           underTest,
           subjectsBuilder.build(),
+          new SerializationContext(dependencies),
+          new DeserializationContext(dependencies),
           skipBadDataTest,
           verificationFunction);
     }