TestThriftToDynamicProto.java

package com.twitter.elephantbird.util;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Set;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.protobuf.Descriptors.DescriptorValidationException;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.Descriptors.FileDescriptor;
import com.google.protobuf.Message;

import com.twitter.elephantbird.thrift.test.AddressBook;
import com.twitter.elephantbird.thrift.test.MapStruct;
import com.twitter.elephantbird.thrift.test.Name;
import com.twitter.elephantbird.thrift.test.Person;
import com.twitter.elephantbird.thrift.test.PhoneNumber;
import com.twitter.elephantbird.thrift.test.PhoneType;
import com.twitter.elephantbird.thrift.test.PrimitiveListsStruct;
import com.twitter.elephantbird.thrift.test.PrimitiveSetsStruct;
import com.twitter.elephantbird.util.Protobufs;

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

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

public class TestThriftToDynamicProto {

  @Rule
  public ExpectedException exception = ExpectedException.none();

  private PhoneNumber genPhoneNumber(String number, PhoneType type) {
    PhoneNumber phoneNumber = new PhoneNumber(number);
    phoneNumber.setType(type);
    return phoneNumber;
  }

  private List<PhoneNumber> genPhoneList() {
    List<PhoneNumber> phoneList = Lists.newLinkedList();
    phoneList.add(genPhoneNumber("123", PhoneType.HOME));
    phoneList.add(genPhoneNumber("456", PhoneType.MOBILE));
    phoneList.add(genPhoneNumber("789", PhoneType.WORK));
    return phoneList;
  }

  private Person genPerson() {
    Person person = new Person();
    person.setName(new Name("Johnny", "Appleseed"));
    person.setId(123);
    person.setEmail("email@address.com");
    person.setPhones(genPhoneList());
    return person;
  }

  private PrimitiveListsStruct genPrimitiveListsStruct() {
    PrimitiveListsStruct listStruct = new PrimitiveListsStruct();

    List<String> strings = Lists.newLinkedList();
    strings.add("string1");
    strings.add("string2");
    strings.add("string3");

    List<Long> longs = Lists.newLinkedList();
    longs.add(1L);
    longs.add(2L);
    longs.add(3L);

    return listStruct.setStrings(strings).setLongs(longs);
  }

  private PrimitiveSetsStruct genPrimitiveSetsStruct() {
    PrimitiveSetsStruct setsStruct = new PrimitiveSetsStruct();

    Set<String> strings = Sets.newHashSet();
    strings.add("string1");
    strings.add("string2");
    strings.add("string3");

    Set<Long> longs = Sets.newHashSet();
    longs.add(1L);
    longs.add(2L);
    longs.add(3L);

    return setsStruct.setStrings(strings).setLongs(longs);
  }

  private MapStruct genMapStruct() {
    Map<Integer, String> map = Maps.newHashMap();
    map.put(1, "one");
    map.put(2, "two");
    map.put(3, "three");
    return new MapStruct(map);
  }

  private void compareName(Name expected, Message actual) {
    assertNotNull(actual);
    assertEquals(expected.first_name, Protobufs.getFieldByName(actual, "first_name"));
    assertEquals(expected.last_name,  Protobufs.getFieldByName(actual, "last_name"));
  }

  private void comparePhoneNumbers(List<PhoneNumber> expected, List<?> actual) {
    assertNotNull(actual);
    assertEquals(expected.size(), actual.size());
    for (int i = 0; i < expected.size(); i++) {
      PhoneNumber number = expected.get(i);
      Message msg = (Message) actual.get(i);
      assertEquals(number.number, Protobufs.getFieldByName(msg, "number"));
      assertEquals(number.type.toString(), Protobufs.getFieldByName(msg, "type"));
    }
  }

  private void comparePrimitiveListStruct(PrimitiveListsStruct expected, Message actual) {
    List<String> strings = (List<String>) Protobufs.getFieldByName(actual, "strings");
    List<Long> longs = (List<Long>) Protobufs.getFieldByName(actual, "longs");
    assertEquals(expected.getStrings(), strings);
    assertEquals(expected.getLongs(), longs);
  }

  // We iterate through the Message and confirm its List elements match the Set elements
  private void comparePrimitiveSetsStruct(PrimitiveSetsStruct expected, Message actual) {
    List<String> actualStrings = (List<String>) Protobufs.getFieldByName(actual, "strings");
    List<Long> actualLongs = (List<Long>) Protobufs.getFieldByName(actual, "longs");

    Set<String> expectedStrings = expected.strings;
    Set<Long> expectedLongs = expected.longs;

    for (String entry : actualStrings) {
      assertTrue(expectedStrings.remove(entry));
    }
    assertEquals(expectedStrings.size(), 0);

    for (Long entry : actualLongs) {
      assertTrue(expectedLongs.remove(entry));
    }
    assertEquals(expectedLongs.size(), 0);
  }

  @Test
  public void testSimpleStructConversion() throws DescriptorValidationException {
    PhoneNumber thriftPhone = genPhoneNumber("123-34-5467", PhoneType.HOME);
    ThriftToDynamicProto<PhoneNumber> thriftToProto = new ThriftToDynamicProto<PhoneNumber>(PhoneNumber.class);
    Message msg = thriftToProto.convert(thriftPhone);
    assertEquals(thriftPhone.number, Protobufs.getFieldByName(msg, "number"));
    assertEquals(thriftPhone.type.toString(), Protobufs.getFieldByName(msg, "type"));
  }

  @Test
  public void testNestedStructsWhenEnabled() throws DescriptorValidationException {
    ThriftToDynamicProto<Person> thriftToProto =
      new ThriftToDynamicProto<Person>(Person.class, true, false);
    Person person = genPerson();
    Message msg = thriftToProto.convert(person);
    assertEquals(person.id, Protobufs.getFieldByName(msg, "id"));
    assertEquals(person.email, Protobufs.getFieldByName(msg, "email"));
    compareName(person.name, (Message) Protobufs.getFieldByName(msg, "name"));
    comparePhoneNumbers(person.phones, (List) Protobufs.getFieldByName(msg, "phones"));
  }

  @Test
  public void testNestedStructsWhenDisabled() throws DescriptorValidationException {
    ThriftToDynamicProto<Person> thriftToProto =
      new ThriftToDynamicProto<Person>(Person.class);
    Person person = genPerson();
    Message msg = thriftToProto.convert(person);
    assertEquals(person.id, Protobufs.getFieldByName(msg, "id"));
    assertEquals(person.email, Protobufs.getFieldByName(msg, "email"));
    // nested structs not converted
    assertTrue(!Protobufs.hasFieldByName(msg, "name"));
    assertTrue(!Protobufs.hasFieldByName(msg, "phones"));
  }

  @Test
  public void testPrimitiveListConversionWhenNestedStructsEnabled() throws DescriptorValidationException {
    ThriftToDynamicProto<PrimitiveListsStruct> thriftToProto =
      new ThriftToDynamicProto<PrimitiveListsStruct>(PrimitiveListsStruct.class, true, false);
    PrimitiveListsStruct listStruct = genPrimitiveListsStruct();
    Message msg = thriftToProto.convert(listStruct);
    comparePrimitiveListStruct(listStruct, msg);
  }

  @Test
  public void testPrimitiveListConversionWhenNestedStructsDisabled() throws DescriptorValidationException {
    ThriftToDynamicProto<PrimitiveListsStruct> thriftToProto =
      new ThriftToDynamicProto<PrimitiveListsStruct>(PrimitiveListsStruct.class);
    PrimitiveListsStruct listStruct = genPrimitiveListsStruct();
    Message msg = thriftToProto.convert(listStruct);
    comparePrimitiveListStruct(listStruct, msg);
  }

  @Test
  public void testPrimitiveSetConversionWhenNestedStructsDisabled() throws DescriptorValidationException {
    ThriftToDynamicProto<PrimitiveSetsStruct> thriftToProto =
      new ThriftToDynamicProto<PrimitiveSetsStruct>(PrimitiveSetsStruct.class);
    PrimitiveSetsStruct setsStruct = genPrimitiveSetsStruct();
    Message msg = thriftToProto.convert(setsStruct);
    comparePrimitiveSetsStruct(setsStruct, msg);
  }

  @Test
  public void testPrimitiveSetConversionWhenNestedStructsEnabled() throws DescriptorValidationException {
    ThriftToDynamicProto<PrimitiveSetsStruct> thriftToProto =
      new ThriftToDynamicProto<PrimitiveSetsStruct>(PrimitiveSetsStruct.class, true, false);
    PrimitiveSetsStruct setsStruct = genPrimitiveSetsStruct();
    Message msg = thriftToProto.convert(setsStruct);
    comparePrimitiveSetsStruct(setsStruct, msg);
  }

  @Test
  public void testMapConversionWhenNestedStructsEnabled() throws DescriptorValidationException {
    ThriftToDynamicProto<MapStruct> thriftToProto =
      new ThriftToDynamicProto<MapStruct>(MapStruct.class, true, false);
    MapStruct mapStruct = genMapStruct();
    Message msg = thriftToProto.convert(mapStruct);

    Map<Integer, String> expected = mapStruct.entries;
    List<?> entries = (List)Protobufs.getFieldByName(msg, "entries");
    Set<?> expectedKeys = Sets.newHashSet(expected.keySet());
    for (Object entry : entries) {
      Message entryMsg = (Message) entry;
      Object key = Protobufs.getFieldByName(entryMsg, "key");
      assertTrue(expectedKeys.remove(key));

      Object value = Protobufs.getFieldByName(entryMsg, "value");
      assertEquals(expected.get(key), value);
    }

    assertEquals(0, expectedKeys.size());
  }

  @Test
  public void testMapConversionWhenNestedStructsDisabled() throws DescriptorValidationException {
    ThriftToDynamicProto<MapStruct> thriftToProto =
      new ThriftToDynamicProto<MapStruct>(MapStruct.class);
    MapStruct mapStruct = genMapStruct();
    Message msg = thriftToProto.convert(mapStruct);
    assertTrue(!Protobufs.hasFieldByName(msg, "entries"));
  }

  @Test
  public void testBadThriftTypeForGetFieldDescriptor() throws DescriptorValidationException {
    ThriftToDynamicProto<PhoneNumber> converter = new ThriftToDynamicProto<PhoneNumber>(PhoneNumber.class);

    exception.expect(IllegalStateException.class);
    converter.getFieldDescriptor(Person.class, "some_field");
  }

  @Test
  public void testGetFieldTypeDescriptor() throws DescriptorValidationException {
    ThriftToDynamicProto<Person> converter = new ThriftToDynamicProto<Person>(Person.class);
    Person person = genPerson();
    Message msg = converter.convert(person);

    FieldDescriptor expectedFd = msg.getDescriptorForType().findFieldByName("email");
    FieldDescriptor actualFd = converter.getFieldDescriptor(Person.class, "email");

    assertEquals(expectedFd, actualFd);
  }

  @Test
  public void testGetFileDescriptor() throws DescriptorValidationException {
    ThriftToDynamicProto<Person> converter = new ThriftToDynamicProto<Person>(Person.class);
    Person person = genPerson();
    Message msg = converter.convert(person);

    FileDescriptor expectedFd = msg.getDescriptorForType().getFile();
    FileDescriptor actualFd = converter.getFileDescriptor();

    assertEquals(expectedFd, actualFd);
  }
}