test: move common logic into TestBase

substrait-io · Nov 2, 2023 · bea6074 · bea6074
1 parent 33ca926
commit bea6074
Show file tree

Hide file tree

Showing 9 changed files with 29 additions and 71 deletions.
diff --git a/core/src/test/java/io/substrait/TestBase.java b/core/src/test/java/io/substrait/TestBase.java
@@ -1,6 +1,14 @@
 package io.substrait;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import io.substrait.dsl.SubstraitBuilder;
+import io.substrait.extension.ExtensionCollector;
 import io.substrait.extension.SimpleExtension;
+import io.substrait.relation.ProtoRelConverter;
+import io.substrait.relation.Rel;
+import io.substrait.relation.RelProtoConverter;
+import io.substrait.type.TypeCreator;
 import java.io.IOException;
 
 public abstract class TestBase {
@@ -14,4 +22,18 @@ public abstract class TestBase {
       throw new RuntimeException(e);
     }
   }
+
+  protected TypeCreator R = TypeCreator.REQUIRED;
+
+  protected SubstraitBuilder b = new SubstraitBuilder(defaultExtensionCollection);
+  protected ExtensionCollector functionCollector = new ExtensionCollector();
+  protected RelProtoConverter relProtoConverter = new RelProtoConverter(functionCollector);
+  protected ProtoRelConverter protoRelConverter =
+      new ProtoRelConverter(functionCollector, defaultExtensionCollection);
+
+  protected void verifyRoundTrip(Rel rel) {
+    io.substrait.proto.Rel protoRel = relProtoConverter.toProto(rel);
+    Rel relReturned = protoRelConverter.from(protoRel);
+    assertEquals(rel, relReturned);
+  }
 }
diff --git a/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java b/core/src/test/java/io/substrait/relation/ProtoRelConverterTest.java
@@ -15,12 +15,6 @@
 
 public class ProtoRelConverterTest extends TestBase {
 
-  final SubstraitBuilder b = new SubstraitBuilder(defaultExtensionCollection);
-  final ExtensionCollector functionCollector = new ExtensionCollector();
-  final RelProtoConverter relProtoConverter = new RelProtoConverter(functionCollector);
-  final ProtoRelConverter protoRelConverter =
-      new ProtoRelConverter(functionCollector, defaultExtensionCollection);
-
   final NamedScan commonTable =
       b.namedScan(Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
 

diff --git a/core/src/test/java/io/substrait/type/proto/AggregateRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/AggregateRoundtripTest.java
@@ -21,8 +21,6 @@
 import org.junit.jupiter.api.Test;
 
 public class AggregateRoundtripTest extends TestBase {
-  static final org.slf4j.Logger logger =
-      org.slf4j.LoggerFactory.getLogger(AggregateRoundtripTest.class);
 
   private void assertAggregateRoundtrip(Expression.AggregationInvocation invocation) {
     var expression = ExpressionCreator.decimal(false, BigDecimal.TEN, 10, 2);

diff --git a/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/ExtensionRoundtripTest.java
@@ -3,11 +3,8 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 import io.substrait.TestBase;
-import io.substrait.dsl.SubstraitBuilder;
 import io.substrait.expression.Expression;
 import io.substrait.extension.AdvancedExtension;
-import io.substrait.extension.ExtensionCollector;
-import io.substrait.extension.SimpleExtension;
 import io.substrait.relation.Aggregate;
 import io.substrait.relation.Cross;
 import io.substrait.relation.ExtensionLeaf;
@@ -22,7 +19,6 @@
 import io.substrait.relation.Project;
 import io.substrait.relation.ProtoRelConverter;
 import io.substrait.relation.Rel;
-import io.substrait.relation.RelProtoConverter;
 import io.substrait.relation.Set;
 import io.substrait.relation.Sort;
 import io.substrait.relation.VirtualTableScan;
@@ -45,15 +41,8 @@
  */
 public class ExtensionRoundtripTest extends TestBase {
 
-  TypeCreator R = TypeCreator.REQUIRED;
-
-  final SimpleExtension.ExtensionCollection extensions = defaultExtensionCollection;
-
-  final SubstraitBuilder b = new SubstraitBuilder(extensions);
-  final ExtensionCollector functionCollector = new ExtensionCollector();
-  final RelProtoConverter relProtoConverter = new RelProtoConverter(functionCollector);
   final ProtoRelConverter protoRelConverter =
-      new StringHolderHandlingProtoRelConverter(functionCollector, extensions);
+      new StringHolderHandlingProtoRelConverter(functionCollector, defaultExtensionCollection);
 
   final Rel commonTable =
       b.namedScan(Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
@@ -72,7 +61,8 @@ public class ExtensionRoundtripTest extends TestBase {
           .optimization(new StringHolder("REL OPTIMIZATION"))
           .build();
 
-  void verifyRoundTrip(Rel rel) {
+  @Override
+  protected void verifyRoundTrip(Rel rel) {
     io.substrait.proto.Rel protoRel = relProtoConverter.toProto(rel);
     Rel relReturned = protoRelConverter.from(protoRel);
     assertEquals(rel, relReturned);

diff --git a/core/src/test/java/io/substrait/type/proto/GenericRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/GenericRoundtripTest.java
@@ -29,8 +29,6 @@
 import org.junit.jupiter.params.provider.MethodSource;
 
 public class GenericRoundtripTest extends TestBase {
-  static final org.slf4j.Logger logger =
-      org.slf4j.LoggerFactory.getLogger(GenericRoundtripTest.class);
 
   static Random rand = new Random(123);
 

diff --git a/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/JoinRoundtripTest.java
@@ -1,34 +1,14 @@
 package io.substrait.type.proto;
 
-import static org.junit.jupiter.api.Assertions.assertEquals;
-
 import io.substrait.TestBase;
-import io.substrait.dsl.SubstraitBuilder;
-import io.substrait.extension.ExtensionCollector;
-import io.substrait.extension.SimpleExtension;
-import io.substrait.relation.ProtoRelConverter;
 import io.substrait.relation.Rel;
-import io.substrait.relation.RelProtoConverter;
 import io.substrait.relation.physical.HashJoin;
-import io.substrait.relation.utils.StringHolderHandlingProtoRelConverter;
-import io.substrait.type.TypeCreator;
 import java.util.Arrays;
 import java.util.List;
 import org.junit.jupiter.api.Test;
 
 public class JoinRoundtripTest extends TestBase {
 
-  final SimpleExtension.ExtensionCollection extensions = defaultExtensionCollection;
-
-  TypeCreator R = TypeCreator.REQUIRED;
-
-  final SubstraitBuilder b = new SubstraitBuilder(extensions);
-
-  final ExtensionCollector functionCollector = new ExtensionCollector();
-  final RelProtoConverter relProtoConverter = new RelProtoConverter(functionCollector);
-  final ProtoRelConverter protoRelConverter =
-      new StringHolderHandlingProtoRelConverter(functionCollector, extensions);
-
   final Rel leftTable =
       b.namedScan(
           Arrays.asList("T1"),
@@ -41,12 +21,6 @@ public class JoinRoundtripTest extends TestBase {
           Arrays.asList("d", "e", "f"),
           Arrays.asList(R.FP64, R.STRING, R.I64));
 
-  void verifyRoundTrip(Rel rel) {
-    io.substrait.proto.Rel protoRel = relProtoConverter.toProto(rel);
-    Rel relReturned = protoRelConverter.from(protoRel);
-    assertEquals(rel, relReturned);
-  }
-
   @Test
   void hashJoin() {
     List<Integer> leftKeys = Arrays.asList(0, 1);

diff --git a/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/LiteralRoundtripTest.java
@@ -7,25 +7,16 @@
 import io.substrait.expression.ExpressionCreator;
 import io.substrait.expression.proto.ExpressionProtoConverter;
 import io.substrait.expression.proto.ProtoExpressionConverter;
-import io.substrait.extension.ExtensionCollector;
-import io.substrait.relation.ProtoRelConverter;
 import java.math.BigDecimal;
 import org.junit.jupiter.api.Test;
 
 public class LiteralRoundtripTest extends TestBase {
-  static final org.slf4j.Logger logger =
-      org.slf4j.LoggerFactory.getLogger(LiteralRoundtripTest.class);
 
   @Test
   void decimal() {
     var val = ExpressionCreator.decimal(false, BigDecimal.TEN, 10, 2);
     var to = new ExpressionProtoConverter(null, null);
-    var from =
-        new ProtoExpressionConverter(
-            null,
-            null,
-            EMPTY_TYPE,
-            new ProtoRelConverter(new ExtensionCollector(), defaultExtensionCollection));
+    var from = new ProtoExpressionConverter(null, null, EMPTY_TYPE, protoRelConverter);
     assertEquals(val, from.from(val.accept(to)));
   }
 }
diff --git a/core/src/test/java/io/substrait/type/proto/LocalFilesRoundtripTest.java b/core/src/test/java/io/substrait/type/proto/LocalFilesRoundtripTest.java
@@ -25,15 +25,7 @@
 
 public class LocalFilesRoundtripTest extends TestBase {
 
-  SimpleExtension.ExtensionCollection extensions = defaultExtensionCollection;
-
-  public LocalFilesRoundtripTest() throws IOException {}
-
   private void assertLocalFilesRoundtrip(FileOrFiles file) {
-    ExtensionCollector functionCollector = new ExtensionCollector();
-    RelProtoConverter to = new RelProtoConverter(functionCollector);
-    ProtoRelConverter from = new ProtoRelConverter(functionCollector, extensions);
-
     var builder =
         LocalFiles.builder()
             .initialSchema(
@@ -47,7 +39,7 @@ private void assertLocalFilesRoundtrip(FileOrFiles file) {
                     .build())
             .addItems(file);
 
-    extensions.scalarFunctions().stream()
+      defaultExtensionCollection.scalarFunctions().stream()
         .filter(s -> s.name().equalsIgnoreCase("equal"))
         .findFirst()
         .map(
@@ -63,9 +55,9 @@ private void assertLocalFilesRoundtrip(FileOrFiles file) {
         .ifPresent(builder::filter);
 
     var localFiles = builder.build();
-    var protoFileRel = to.toProto(localFiles);
+    var protoFileRel = relProtoConverter.toProto(localFiles);
     assertTrue(protoFileRel.getRead().hasFilter());
-    assertEquals(protoFileRel, to.toProto(from.from(protoFileRel)));
+    assertEquals(protoFileRel, relProtoConverter.toProto(protoRelConverter.from(protoFileRel)));
   }
 
   private ImmutableFileOrFiles.Builder setPath(

diff --git a/core/src/test/java/io/substrait/type/proto/TestTypeRoundtrip.java b/core/src/test/java/io/substrait/type/proto/TestTypeRoundtrip.java
@@ -10,7 +10,6 @@
 import org.junit.jupiter.params.provider.ValueSource;
 
 public class TestTypeRoundtrip {
-  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(TestTypeRoundtrip.class);
 
   @ParameterizedTest
   @ValueSource(booleans = {true, false})