Merge pull request #6311 from pabloem/twsit

[BEAM-5377] Adding TopWikipediaSessions integration test

examples/java/src/main/java/org/apache/beam/examples/complete/TopWikipediaSessions.java

@@ -19,6 +19,7 @@
 
 import com.google.api.services.bigquery.model.TableRow;
 import java.io.IOException;
+import java.math.BigDecimal;
 import java.util.List;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.io.TextIO;
@@ -76,7 +77,13 @@ static class ExtractUserAndTimestamp extends DoFn<TableRow, String> {
     @ProcessElement
     public void processElement(ProcessContext c) {
       TableRow row = c.element();
-      int timestamp = (Integer) row.get("timestamp");
+      int timestamp;
+      // TODO(BEAM-5390): Avoid this workaround.
+      try {
+        timestamp = ((BigDecimal) row.get("timestamp")).intValue();
+      } catch (ClassCastException e) {
+        timestamp = ((Integer) row.get("timestamp")).intValue();
+      }
       String userName = (String) row.get("contributor_username");
       if (userName != null) {
         // Sets the implicit timestamp field to be used in windowing.
@@ -191,18 +198,21 @@ public interface Options extends PipelineOptions {
     void setOutput(String value);
   }
 
-  public static void main(String[] args) {
-    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
-
+  public static void run(Options options) {
     Pipeline p = Pipeline.create(options);
 
     double samplingThreshold = 0.1;
 
     p.apply(TextIO.read().from(options.getInput()))
         .apply(MapElements.via(new ParseTableRowJson()))
         .apply(new ComputeTopSessions(samplingThreshold))
-        .apply("Write", TextIO.write().withoutSharding().to(options.getOutput()));
+        .apply("Write", TextIO.write().to(options.getOutput()));
 
     p.run().waitUntilFinish();
   }
+
+  public static void main(String[] args) {
+    Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
+    run(options);
+  }
 }

examples/java/src/test/java/org/apache/beam/examples/complete/TopWikipediaSessionsIT.java

@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.beam.examples.complete;
+
+import java.util.Date;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.FileChecksumMatcher;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.testing.TestPipelineOptions;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** End-to-end tests of TopWikipediaSessions. */
+@RunWith(JUnit4.class)
+public class TopWikipediaSessionsIT {
+
+  private static final String DEFAULT_INPUT_10_FILES =
+      "gs://apache-beam-samples/wikipedia_edits/wiki_data-00000000000*.json";
+  private static final String DEFAULT_OUTPUT_CHECKSUM = "a7f0c50b895d0a2e37b78c3f94eadcfb11a647a6";
+
+  /** PipelineOptions for the TopWikipediaSessions integration test. */
+  public interface TopWikipediaSessionsITOptions
+      extends TestPipelineOptions, TopWikipediaSessions.Options {}
+
+  @BeforeClass
+  public static void setUp() {
+    PipelineOptionsFactory.register(TestPipelineOptions.class);
+  }
+
+  @Test
+  public void testE2ETopWikiPages() throws Exception {
+    TopWikipediaSessionsITOptions options =
+        TestPipeline.testingPipelineOptions().as(TopWikipediaSessionsITOptions.class);
+
+    options.setInput(DEFAULT_INPUT_10_FILES);
+    options.setOutput(
+        FileSystems.matchNewResource(options.getTempRoot(), true)
+            .resolve(
+                String.format("topwikisessions-it-%tF-%<tH-%<tM-%<tS-%<tL", new Date()),
+                StandardResolveOptions.RESOLVE_DIRECTORY)
+            .resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY)
+            .resolve("results", StandardResolveOptions.RESOLVE_FILE)
+            .toString());
+    options.setOnSuccessMatcher(
+        new FileChecksumMatcher(DEFAULT_OUTPUT_CHECKSUM, options.getOutput() + "*-of-*"));
+
+    TopWikipediaSessions.run(options);
+  }
+}