Parametrized test

lgajowy · Nov 3, 2017 · e9b7418 · e9b7418
1 parent 6d04d78
commit e9b7418
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/sdks/java/io/text/src/test/java/org/apache/beam/sdk/io/text/IOTextTestPipelineOptions.java b/sdks/java/io/text/src/test/java/org/apache/beam/sdk/io/text/IOTextTestPipelineOptions.java
@@ -0,0 +1,16 @@
+package org.apache.beam.sdk.io.text;
+
+import org.apache.beam.sdk.options.Default;
+import org.apache.beam.sdk.options.Description;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+
+/**
+ * Pipeline options specific for filebased IO tests.
+ */
+
+public interface IOTextTestPipelineOptions extends TestPipelineOptions{
+    @Description("Test row count")
+    @Default.Long(10000)
+    Long getTestSize();
+    void setTestSize(Long value);
+}
diff --git a/sdks/java/io/text/src/test/java/org/apache/beam/sdk/io/text/TextIOIT.java b/sdks/java/io/text/src/test/java/org/apache/beam/sdk/io/text/TextIOIT.java
@@ -2,9 +2,12 @@
 
 import com.google.common.hash.Hashing;
 import java.nio.charset.StandardCharsets;
+import java.text.ParseException;
+
 import org.apache.beam.sdk.io.GenerateSequence;
 import org.apache.beam.sdk.io.TextIO;
 import org.apache.beam.sdk.io.common.HashingFn;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.testing.PAssert;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Combine;
@@ -13,6 +16,7 @@
 import org.apache.beam.sdk.transforms.Values;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.TypeDescriptors;
+import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -23,16 +27,27 @@
 public class TextIOIT {
 
   private static final String FILE_BASENAME = "textioit";
-  private static final long LINES_OF_TEXT_COUNT = 10000L;
+  private static long linesOfTextCount = 10000L;
   private static final String EXPECTED_HASHCODE = "ccae48ff685c1822e9f4d510363bf018";
 
   @Rule public TestPipeline pipeline = TestPipeline.create();
 
+  @BeforeClass
+  public static void setup() throws ParseException {
+    PipelineOptionsFactory.register(IOTextTestPipelineOptions.class);
+    IOTextTestPipelineOptions options = TestPipeline.testingPipelineOptions()
+            .as(IOTextTestPipelineOptions.class);
+
+    if (options.getTestSize() != null) {
+      linesOfTextCount = options.getTestSize();
+    }
+  }
+
   @Test
   public void testWriteThenRead() {
 
     PCollection<String> consolidatedContentHashcode = pipeline
-        .apply("Generate sequence", GenerateSequence.from(0).to(LINES_OF_TEXT_COUNT))
+        .apply("Generate sequence", GenerateSequence.from(0).to(linesOfTextCount))
         .apply("Produce text", MapElements.into(TypeDescriptors.strings()).via(produceTextLine()))
         .apply("Write content to files", TextIO.write().to(FILE_BASENAME).withOutputFilenames())
         .getPerDestinationOutputFilenames()