schemaName provided to creator (#8834)

Consensys · Nov 19, 2024 · 46adf0c · 46adf0c
1 parent c4b29fc
commit 46adf0c
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 145 deletions.
diff --git a/ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/BaseSchemaProvider.java b/ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/BaseSchemaProvider.java
@@ -20,7 +20,6 @@
 import java.util.List;
 import java.util.Set;
 import java.util.TreeMap;
-import java.util.function.BiFunction;
 import tech.pegasys.teku.spec.SpecMilestone;
 import tech.pegasys.teku.spec.config.SpecConfig;
 import tech.pegasys.teku.spec.schemas.registry.SchemaTypes.SchemaId;
@@ -78,7 +77,9 @@ protected T createSchema(
       final SchemaRegistry registry,
       final SpecMilestone effectiveMilestone,
       final SpecConfig specConfig) {
-    return getSchemaCreator(effectiveMilestone).creator.apply(registry, specConfig);
+    return getSchemaCreator(effectiveMilestone)
+        .creator
+        .create(registry, specConfig, schemaId.getSchemaName(registry.getMilestone()));
   }
 
   private SchemaProviderCreator<T> getSchemaCreator(final SpecMilestone milestone) {
@@ -95,8 +96,7 @@ public Set<SpecMilestone> getSupportedMilestones() {
     return milestoneToSchemaCreator.keySet();
   }
 
-  protected record SchemaProviderCreator<T>(
-      SpecMilestone baseMilestone, BiFunction<SchemaRegistry, SpecConfig, T> creator) {
+  protected record SchemaProviderCreator<T>(SpecMilestone baseMilestone, SchemaCreator<T> creator) {
 
     @Override
     public String toString() {
@@ -156,8 +156,7 @@ private Builder(final SchemaId<T> schemaId) {
     }
 
     public Builder<T> withCreator(
-        final SpecMilestone milestone,
-        final BiFunction<SchemaRegistry, SpecConfig, T> creationSchema) {
+        final SpecMilestone milestone, final SchemaCreator<T> creationSchema) {
       checkArgument(
           schemaProviderCreators.isEmpty()
               || milestone.isGreaterThan(schemaProviderCreators.getLast().baseMilestone),
@@ -199,4 +198,9 @@ public BaseSchemaProvider<T> build() {
           schemaId, schemaProviderCreators, untilMilestone, alwaysCreateNewSchema);
     }
   }
+
+  @FunctionalInterface
+  public interface SchemaCreator<T> {
+    T create(SchemaRegistry registry, SpecConfig specConfig, String schemaName);
+  }
 }
diff --git a/ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/SchemaRegistry.java b/ethereum/spec/src/main/java/tech/pegasys/teku/spec/schemas/registry/SchemaRegistry.java
@@ -67,7 +67,7 @@ boolean isProviderRegistered(final SchemaProvider<?> provider) {
 
   @SuppressWarnings("unchecked")
   public <T> T get(final SchemaId<T> schemaId) {
-    SchemaProvider<T> provider = (SchemaProvider<T>) providers.get(schemaId);
+    final SchemaProvider<T> provider = (SchemaProvider<T>) providers.get(schemaId);
     if (provider == null) {
       throw new IllegalArgumentException(
           "No provider registered for schema "