package com.squareup.cascading2.function;

import cascading.flow.hadoop.HadoopFlowProcess;
import cascading.operation.Function;
import cascading.operation.FunctionCall;
import cascading.pipe.Pipe;
import cascading.tuple.Fields;
import cascading.tuple.Tuple;
import cascading.tuple.TupleEntry;
import cascading.tuple.TupleEntryCollector;
import com.squareup.cascading2.generated.Example;
import java.io.IOException;
import java.util.concurrent.atomic.AtomicReference;
import junit.framework.TestCase;

public class ExtractProtoTest extends TestCase {
  private static final Example.Partnership P1 = Example.Partnership.newBuilder()
      .setFollower(Example.Person.newBuilder()
          .setName("Andy")
          .setEmail("andy@")
      )
      .setLeader(Example.Person.newBuilder()
          .setName("Bryan")
          .setEmail("bryan@")
      )
      .build();

  private static final Example.Partnership P2 = Example.Partnership.newBuilder()
      .setLeader(Example.Person.newBuilder().setName("Bryan").setEmail("bryan@"))
      .build();

  private static final Example.Partnership P3 = Example.Partnership.newBuilder()
      .setLeader(Example.Person.newBuilder()
          .setName("Jack")
          .setEmail("jack@")
          .setPosition(Example.Person.Position.CEO))
      .build();

  public void testValidatesNonExistingField() throws Exception {
    try {
      new ExtractProto(Example.Partnership.class, "follower.has_mind_powers");
      fail("Expected an exception!");
    } catch (IllegalArgumentException e) {
      // yay!
    }
  }

  public void testThrowsOnRepeatedField() throws Exception {
    try {
      new ExtractProto(Example.Partnership.class, "silent.name");
      fail("Expected an exception!");
    } catch (IllegalArgumentException e) {
      // yay!
    }
    try {
      new ExtractProto(Example.Partnership.class, "silent");
      fail("Expected an exception!");
    } catch (IllegalArgumentException e) {
      // yay!
    }
  }

  public void testNullFirstObject() throws Exception {
    assertEquals(new Tuple(null, null, null, null),
        exec(new ExtractProto(Example.Partnership.class, "follower.name", "follower.email", "leader.name", "leader.email"), new Tuple((Object)null)));
  }

  public void testAllPresent() throws Exception {
    assertEquals(new Tuple("Andy", "andy@", "Bryan", "bryan@"),
        exec(new ExtractProto(Example.Partnership.class, "follower.name", "follower.email", "leader.name", "leader.email"), new Tuple(P1)));
  }

  public void testSomePresent() throws Exception {
    assertEquals(new Tuple(null, null, "Bryan", "bryan@"),
        exec(new ExtractProto(Example.Partnership.class, "follower.name", "follower.email", "leader.name", "leader.email"), new Tuple(P2)));
  }

  public void testEnum() throws Exception {
    assertEquals(new Tuple("Jack", "jack@", Example.Person.Position.CEO.getNumber()),
        exec(new ExtractProto(Example.Partnership.class, "leader.name", "leader.email", "leader.position"), new Tuple(P3)));
  }

  private static Tuple exec(Function f, final Tuple input) {
    final AtomicReference<Tuple> output = new AtomicReference<Tuple>();
    f.prepare(new HadoopFlowProcess(), null);

    f.operate(new HadoopFlowProcess(), new FunctionCall() {
      @Override public TupleEntry getArguments() {
        return new TupleEntry(new Fields("blah"), input);
      }

      @Override public Fields getDeclaredFields() {
        return null;
      }

      @Override public TupleEntryCollector getOutputCollector() {
        return new TupleEntryCollector() {
          @Override protected void collect(TupleEntry tupleEntry) throws IOException {
            output.set(tupleEntry.getTuple());
          }
        };
      }

      @Override public Object getContext() {
        return null;
      }

      @Override public void setContext(Object o) {
      }

      @Override public Fields getArgumentFields() {
        return null;
      }
    });
    return output.get();
  }
}

@@ -36,16 +36,97 @@
 ction;%0A%0A
+import cascading.flow.FlowDef;%0Aimport cascading.flow.hadoop.HadoopFlowConnector;%0A
 import c
@@ -269,12 +269,73 @@
 ipe.
-Pipe
+Each;%0Aimport cascading.pipe.Pipe;%0Aimport cascading.tap.hadoop.Hfs
 ;%0Aim
@@ -518,24 +518,78 @@
 ed.Example;%0A
+import com.squareup.cascading2.scheme.ProtobufScheme;%0A
 import java.
@@ -688,16 +688,135 @@
 estCase;
+%0Aimport org.apache.hadoop.conf.Configuration;%0Aimport org.apache.hadoop.fs.FileSystem;%0Aimport org.apache.hadoop.fs.Path;
 %0A%0Apublic
@@ -3324,16 +3324,926 @@
 );%0A  %7D%0A%0A
+  public void testSerialization() throws Exception %7B%0A    Pipe p = new Each(%22input%22, new ExtractProto(Example.Partnership.class, %22leader.name%22));%0A%0A    final String INPUT_PATH = %22/tmp/ExtractProtoTest/input%22;%0A    final String OUTPUT_PATH = %22/tmp/ExtractProtoTest/output%22;%0A%0A    Hfs inputTap = new Hfs(new ProtobufScheme(%22partnership%22, Example.Partnership.class),%0A        INPUT_PATH);%0A    Hfs outputTap = new Hfs(new ProtobufScheme(%22leader.name%22, Example.Partnership.class),%0A        OUTPUT_PATH);%0A%0A    // make sure the input path exists%0A    inputTap.openForWrite(new HadoopFlowProcess()).close();%0A    // make sure the output path does not exist%0A    FileSystem.get(new Configuration()).delete(new Path(OUTPUT_PATH), true);%0A%0A    FlowDef flowDef = new FlowDef()%0A        .addSource(%22input%22, inputTap)%0A        .addSink(p, outputTap)%0A        .addTail(p);%0A%0A    new HadoopFlowConnector().connect(flowDef).complete();%0A  %7D%0A%0A
   privat

Add test to make sure ExtractProto serializes properly