/*
 * 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.runners.spark.translation;

import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertThat;

import org.apache.beam.runners.direct.DirectRunner;
import org.apache.beam.runners.spark.SparkRunner;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.runners.PipelineRunner;
import org.apache.beam.sdk.values.PCollection;

import com.google.common.base.Charsets;

import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.List;

/**
 * A test for the transforms registered in TransformTranslator.
 * Builds a regular Dataflow pipeline with each of the mapped
 * transforms, and makes sure that they work when the pipeline is
 * executed in Spark.
 */
public class TransformTranslatorTest {
  @Rule public TemporaryFolder tmp = new TemporaryFolder();

  /**
   * Builds a simple pipeline with TextIO.Read and TextIO.Write, runs the pipeline
   * in DirectRunner and on SparkRunner, with the mapped dataflow-to-spark
   * transforms. Finally it makes sure that the results are the same for both runs.
   */
  @Test
  public void testTextIOReadAndWriteTransforms() throws IOException {
    String directOut = runPipeline(DirectRunner.class);
    String sparkOut = runPipeline(SparkRunner.class);

    List<String> directOutput =
        Files.readAllLines(Paths.get(directOut + "-00000-of-00001"), Charsets.UTF_8);

    List<String> sparkOutput =
        Files.readAllLines(Paths.get(sparkOut + "-00000-of-00001"), Charsets.UTF_8);

    // sort output to get a stable result (PCollections are not ordered)
    assertThat(sparkOutput, containsInAnyOrder(directOutput.toArray()));
  }

  private String runPipeline(Class<? extends PipelineRunner<?>> runner) throws IOException {
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setRunner(runner);
    Pipeline p = Pipeline.create(options);
    File outFile = tmp.newFile();
    PCollection<String> lines =  p.apply(TextIO.Read.from("src/test/resources/test_text.txt"));
    lines.apply(TextIO.Write.to(outFile.getAbsolutePath()));
    p.run();
    return outFile.getAbsolutePath();
  }
}

@@ -1380,16 +1380,56 @@
 rsets;%0A%0A
+import org.apache.commons.io.FileUtils;%0A
 import o
@@ -1505,16 +1505,73 @@
 yFolder;
+%0Aimport org.slf4j.Logger;%0Aimport org.slf4j.LoggerFactory;
 %0A%0Aimport
@@ -1629,50 +1629,45 @@
 ava.
-nio.file.Files;%0Aimport java.nio.file.Paths
+util.ArrayList;%0Aimport java.util.List
 ;%0Aim
@@ -1681,20 +1681,29 @@
 va.util.
-List
+regex.Pattern
 ;%0A%0A/**%0A 
@@ -1789,16 +1789,12 @@
 lar 
-Dataflow
+Beam
  pip
@@ -1954,16 +1954,108 @@
 rTest %7B%0A
+  private static final Logger LOG = LoggerFactory.getLogger(TransformTranslatorTest.class);%0A
   @Rule 
@@ -2558,123 +2558,216 @@
     
-List%3CString%3E directOutput =%0A        Files.readAllLines(Paths.get(directOut + %22-00000-of-00001%22), Charsets.UTF_8
+File directOutFile = new File(directOut);%0A    List%3CString%3E directOutput =%0A            readFromOutputFiles(directOutFile.getParentFile(), directOutFile.getName());%0A%0A    File sparkOutFile = new File(sparkOut
 );%0A
-%0A
     
@@ -2805,82 +2805,84 @@
     
-Files.readAllLines(Paths.get(sparkOut + %22-00000-of-00001%22), Charsets.UTF_8
+    readFromOutputFiles(sparkOutFile.getParentFile(), sparkOutFile.getName()
 );%0A%0A
@@ -3507,11 +3507,665 @@
 h();%0A  %7D
+%0A%0A  private List%3CString%3E readFromOutputFiles(File parent, String outPattern) throws IOException %7B%0A    // example pattern: outprefix-00000-of-00001%0A    Pattern pattern = Pattern.compile(String.format(%22%25s-%5B0-9%5D%7B5%7D-of-%5B0-9%5D%7B5%7D%22, outPattern));%0A    List%3CString%3E lines = new ArrayList%3C%3E();%0A    if (parent.exists() && parent.isDirectory()) %7B%0A      //noinspection ConstantConditions%0A      for (File f : parent.listFiles()) %7B%0A        if (pattern.matcher(f.getName()).matches()) %7B%0A          LOG.info(%22For %22 + outPattern + %22 reading file %22 + f.getName());%0A          lines.addAll(FileUtils.readLines(f, Charsets.UTF_8));%0A        %7D%0A      %7D%0A    %7D%0A    return lines;%0A  %7D
 %0A%7D%0A

Read all matching output files since default sharding in DirectRunner is not constant.