From e9b7418539f2aec4449ab480705179fb27be0e85 Mon Sep 17 00:00:00 2001
From: Kamil Szewczyk <szewinho@gmail.com>
Date: Fri, 3 Nov 2017 11:06:21 +0100
Subject: [PATCH] Parametrized test

---
 .../io/text/IOTextTestPipelineOptions.java    | 16 ++++++++++++++++
 .../org/apache/beam/sdk/io/text/TextIOIT.java | 19 +++++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 sdks/java/io/text/src/test/java/org/apache/beam/sdk/io/text/IOTextTestPipelineOptions.java

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
new file mode 100644
index 0000000000000..9a146300f1372
--- /dev/null
+++ 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
index f2c251c260f9c..cabbd4532c029 100644
--- 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()