1 files changed, 80 insertions, 18 deletions

diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
index 71efbd714b..dca168b6c6 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
@@ -39,17 +39,17 @@ public class ObjectCodecRegistryTest {
             .add(codec2)
             .build();
 
-    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor(String.class);
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptorForObject("hello");
     assertThat(fooDescriptor.getCodec()).isSameAs(codec1);
     assertThat(underTest.getCodecDescriptorByTag(fooDescriptor.getTag())).isSameAs(fooDescriptor);
 
-    CodecDescriptor barDescriptor = underTest.getCodecDescriptor(Integer.class);
+    CodecDescriptor barDescriptor = underTest.getCodecDescriptorForObject(1);
     assertThat(barDescriptor.getCodec()).isSameAs(codec2);
     assertThat(underTest.getCodecDescriptorByTag(barDescriptor.getTag())).isSameAs(barDescriptor);
 
     assertThat(barDescriptor.getTag()).isNotEqualTo(fooDescriptor.getTag());
 
-    assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptor(Byte.class));
+    assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorForObject((byte) 1));
     assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
   }
 
@@ -58,18 +58,24 @@ public class ObjectCodecRegistryTest {
     SingletonCodec<String> codec = SingletonCodec.of("value1", "mnemonic1");
 
     ObjectCodecRegistry underTest =
-        ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(true).add(codec).build();
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(codec)
+            .addClassName(Byte.class.getName())
+            .addClassName(Integer.class.getName())
+            .build();
 
-    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor(String.class);
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptorForObject("value1");
     assertThat(fooDescriptor.getCodec()).isSameAs(codec);
 
-    CodecDescriptor barDefaultDescriptor = underTest.getCodecDescriptor(Integer.class);
+    CodecDescriptor barDefaultDescriptor = underTest.getCodecDescriptorForObject(15);
     assertThat(barDefaultDescriptor.getCodec()).isNotSameAs(codec);
     assertThat(barDefaultDescriptor.getTag()).isNotEqualTo(fooDescriptor.getTag());
     assertThat(underTest.getCodecDescriptorByTag(barDefaultDescriptor.getTag()))
         .isSameAs(barDefaultDescriptor);
 
-    assertThat(underTest.getCodecDescriptor(Byte.class)).isSameAs(barDefaultDescriptor);
+    assertThat(underTest.getCodecDescriptorForObject((byte) 9).getCodec().getClass())
+        .isSameAs(barDefaultDescriptor.getCodec().getClass());
 
     // Bogus tags still throw.
     assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
@@ -81,18 +87,28 @@ public class ObjectCodecRegistryTest {
     SingletonCodec<Integer> codec2 = SingletonCodec.of(1, "mnemonic2");
 
     ObjectCodecRegistry underTest1 =
-        ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(true).add(codec1).add(codec2).build();
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(codec1)
+            .add(codec2)
+            .addClassName(Byte.class.getName())
+            .build();
 
     ObjectCodecRegistry underTest2 =
-        ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(true).add(codec2).add(codec1).build();
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(codec2)
+            .add(codec1)
+            .addClassName(Byte.class.getName())
+            .build();
 
-    assertThat(underTest1.getCodecDescriptor(String.class).getTag())
-        .isEqualTo(underTest2.getCodecDescriptor(String.class).getTag());
-    assertThat(underTest1.getCodecDescriptor(Integer.class).getTag())
-        .isEqualTo(underTest2.getCodecDescriptor(Integer.class).getTag());
+    assertThat(underTest1.getCodecDescriptorForObject("value1").getTag())
+        .isEqualTo(underTest2.getCodecDescriptorForObject("value1").getTag());
+    assertThat(underTest1.getCodecDescriptorForObject(5).getTag())
+        .isEqualTo(underTest2.getCodecDescriptorForObject(5).getTag());
     // Default codec.
-    assertThat(underTest1.getCodecDescriptor(Byte.class).getTag())
-        .isEqualTo(underTest2.getCodecDescriptor(Byte.class).getTag());
+    assertThat(underTest1.getCodecDescriptorForObject((byte) 10).getTag())
+        .isEqualTo(underTest2.getCodecDescriptorForObject((byte) 10).getTag());
   }
 
   @Test
@@ -133,9 +149,55 @@ public class ObjectCodecRegistryTest {
             .build();
 
     ObjectCodecRegistry copy = underTest.getBuilder().build();
-    assertThat(copy.getCodecDescriptor(Integer.class).getTag()).isEqualTo(1);
-    assertThat(copy.getCodecDescriptor(String.class).getTag()).isEqualTo(2);
+    assertThat(copy.getCodecDescriptorForObject(12).getTag()).isEqualTo(1);
+    assertThat(copy.getCodecDescriptorForObject("value1").getTag()).isEqualTo(2);
     assertThat(copy.maybeGetTagForConstant(constant)).isNotNull();
-    assertThrows(NoCodecException.class, () -> copy.getCodecDescriptor(Byte.class));
+    assertThrows(NoCodecException.class, () -> copy.getCodecDescriptorForObject((byte) 5));
+  }
+
+  private enum TestEnum {
+    ONE {
+      @Override
+      public int val() {
+        return 1;
+      }
+    },
+    TWO {
+      @Override
+      public int val() {
+        return 2;
+      }
+    },
+    THREE {
+      @Override
+      public int val() {
+        return 3;
+      }
+    };
+
+    public abstract int val();
+  }
+
+  @SuppressWarnings("GetClassOnEnum")
+  @Test
+  public void testDefaultEnum() throws NoCodecException {
+    assertThat(TestEnum.ONE.getClass()).isNotEqualTo(TestEnum.class);
+    assertThat(TestEnum.ONE.getDeclaringClass()).isEqualTo(TestEnum.class);
+    assertThat(TestEnum.ONE.getClass()).isNotEqualTo(TestEnum.TWO.getClass());
+
+    ObjectCodecRegistry underTest =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .addClassName(TestEnum.class.getName())
+            .addClassName(TestEnum.ONE.getClass().getName())
+            .addClassName(TestEnum.TWO.getClass().getName())
+            .addClassName(TestEnum.THREE.getClass().getName())
+            .build();
+
+    CodecDescriptor oneDescriptor = underTest.getCodecDescriptorForObject(TestEnum.ONE);
+    CodecDescriptor twoDescriptor = underTest.getCodecDescriptorForObject(TestEnum.TWO);
+    assertThat(oneDescriptor).isEqualTo(twoDescriptor);
+
+    assertThat(oneDescriptor.getCodec().getEncodedClass()).isEqualTo(TestEnum.class);
   }
 }