Fixtures.java

package com.twitter.elephantbird.pig.piggybank;

import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.data.DataBag;
import org.apache.pig.data.NonSpillableDataBag;
import org.apache.pig.data.Tuple;
import org.apache.pig.data.TupleFactory;

import com.google.common.collect.Lists;

import com.twitter.data.proto.tutorial.AddressBookProtos.AddressBook;
import com.twitter.data.proto.tutorial.AddressBookProtos.Person;
import com.twitter.data.proto.tutorial.AddressBookProtos.Person.PhoneNumber;
import com.twitter.data.proto.tutorial.AddressBookProtos.Person.PhoneType;

public class Fixtures {

  public static TupleFactory tf_ = TupleFactory.getInstance();

  public static AddressBook buildAddressBookProto() {
    AddressBook abProto = AddressBook.newBuilder()
    .addPerson(buildPersonProto())
    .addPerson(buildPersonProto())
    .build();
    return abProto;
  }
  
  public static Person buildPersonProto() {
    return Person.newBuilder()
    .addPhone(makePhoneNumber("415-999-9999"))
    .addPhone(makePhoneNumber("415-666-6666",PhoneType.MOBILE))
    .addPhone(makePhoneNumber("415-333-3333", PhoneType.WORK))
    .setEmail("elephant@bird.com")
    .setName("Elephant Bird")
    .setId(123)
    .build();
  }
  
  public static Tuple buildPersonTuple() throws ExecException {
    DataBag phoneBag = new NonSpillableDataBag(
      Lists.newArrayList(makePhoneNumberTuple("415-999-9999", null), 
      makePhoneNumberTuple("415-666-6666", "MOBILE"),
      makePhoneNumberTuple("415-333-3333", "WORK")));
    
    Tuple entryTuple = tf_.newTuple(4);
    entryTuple.set(0, "Elephant Bird");
    entryTuple.set(1, 123);
    entryTuple.set(2, "elephant@bird.com");
    entryTuple.set(3, phoneBag);
    return entryTuple;
  }
  // {(Elephant Bird,123,elephant@bird.com,{(415-999-9999,HOME),(415-666-6666,MOBILE),(415-333-3333,WORK)})}
  public static Tuple buildAddressBookTuple() throws ExecException {
    DataBag entryBag = new NonSpillableDataBag(
     Lists.newArrayList(buildPersonTuple(),
     buildPersonTuple()));
    return tf_.newTuple(entryBag);
  }

  public static Tuple makePhoneNumberTuple(String number, String type) throws ExecException {
    Tuple t = tf_.newTuple(2);
    t.set(0, number);
    t.set(1, type);
    return t;
  }
  public static PhoneNumber makePhoneNumber(String number) {
    return makePhoneNumber(number, null);
  }
  
  public static PhoneNumber makePhoneNumber(String number, PhoneType type) {
    PhoneNumber.Builder phoneBuilder = PhoneNumber.newBuilder();
    phoneBuilder.setNumber(number);
    if (type != null) {
      phoneBuilder.setType(type);
    }
    return phoneBuilder.build();
  }
}