[BEAM-7690] Port WordCountTest off DoFnTester

kamilwu · Jul 9, 2019 · 4ba9989 · 4ba9989
2 parents 7d7f897 + 51c8873
commit 4ba9989
Showing 1 changed file with 7 additions and 12 deletions.
diff --git a/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java b/examples/java/src/test/java/org/apache/beam/examples/WordCountTest.java
@@ -28,11 +28,9 @@
 import org.apache.beam.sdk.testing.ValidatesRunner;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
-import org.apache.beam.sdk.transforms.DoFnTester;
 import org.apache.beam.sdk.transforms.MapElements;
+import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.PCollection;
-import org.hamcrest.CoreMatchers;
-import org.junit.Assert;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
@@ -46,15 +44,12 @@ public class WordCountTest {
   /** Example test that tests a specific {@link DoFn}. */
   @Test
   public void testExtractWordsFn() throws Exception {
-    DoFnTester<String, String> extractWordsFn = DoFnTester.of(new ExtractWordsFn());
-
-    Assert.assertThat(
-        extractWordsFn.processBundle(" some  input  words "),
-        CoreMatchers.hasItems("some", "input", "words"));
-    Assert.assertThat(extractWordsFn.processBundle(" "), CoreMatchers.hasItems());
-    Assert.assertThat(
-        extractWordsFn.processBundle(" some ", " input", " words"),
-        CoreMatchers.hasItems("some", "input", "words"));
+    List<String> words = Arrays.asList(" some  input  words ", " ", " cool ", " foo", " bar");
+    PCollection<String> output =
+        p.apply(Create.of(words).withCoder(StringUtf8Coder.of()))
+            .apply(ParDo.of(new ExtractWordsFn()));
+    PAssert.that(output).containsInAnyOrder("some", "input", "words", "cool", "foo", "bar");
+    p.run().waitUntilFinish();
   }
 
   static final String[] WORDS_ARRAY =