com.adgear.anoa.test.AnoaTestSample Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of anoa-tests Show documentation
There is a newer version: 3.1.2
package com.adgear.anoa.test;

import com.google.protobuf.ExtensionRegistry;

import com.adgear.anoa.test.ad_exchange.AdExchangeProtobuf;
import com.adgear.anoa.test.ad_exchange.LogEventAvro;
import com.adgear.anoa.test.ad_exchange.LogEventThrift;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.googlecode.protobuf.format.JsonFormat;

import org.apache.avro.Schema;
import org.apache.avro.file.DataFileStream;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.DatumReader;
import org.apache.avro.io.DatumWriter;
import org.apache.avro.io.Decoder;
import org.apache.avro.io.DecoderFactory;
import org.apache.avro.io.Encoder;
import org.apache.avro.io.EncoderFactory;
import org.apache.avro.specific.SpecificDatumReader;
import org.apache.avro.specific.SpecificDatumWriter;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TJSONProtocol;
import org.apache.thrift.transport.TMemoryInputTransport;
import org.jooq.lambda.Unchecked;
import org.jooq.lambda.fi.util.function.CheckedConsumer;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UncheckedIOException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Supplier;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import junitx.framework.ListAssert;

public class AnoaTestSample {

  /* RESOURCES */

  static private final JsonTestResource JSON = new JsonTestResource("/json.json");
  static private final JsonTestResource AVRO = new JsonTestResource("/avro.json");
  static private final JsonTestResource THRIFT = new JsonTestResource("/thrift.json");
  static private final JsonTestResource PROTOBUF = new JsonTestResource("/protobuf.json");


  /* PROPERTIES */

  public final int n = 1000;
  public final long nl = n;

  public final Schema avroSchema = LogEventAvro.getClassSchema();
  public final Class avroClass = LogEventAvro.class;

  protected final List avroPojos = new ArrayList<>();
  protected final List avroBinaries = new ArrayList<>();
  protected final byte[] avroBatch;
  protected final List avroGenerics = new ArrayList<>();

  public final Class thriftClass = LogEventThrift.class;
  public final Supplier thriftSupplier = LogEventThrift::new;

  protected final List thriftPojos = new ArrayList<>();
  protected final List thriftBinaries = new ArrayList<>();
  protected final List thriftCompacts = new ArrayList<>();

  public final Class protobufClass =
      AdExchangeProtobuf.LogEvent.class;
  public final Supplier protobufSupplier =
      AdExchangeProtobuf.LogEvent::newBuilder;

  protected final List protobufPojos = new ArrayList<>();
  protected final List protobufBinaries = new ArrayList<>();


  /* CONSTRUCTOR */

  public AnoaTestSample() {
    try {
      avroBatch = buildAvro();
      buildThrift();
      buildProtobuf();
    } catch (Exception e) {
      throw new IllegalStateException(e);
    }
  }

  private byte[] buildAvro() throws IOException {
    DatumReader specificReader = new SpecificDatumReader<>(avroClass);
    for (String json : AVRO.jsonStrings) {
      Decoder decoder = DecoderFactory.get().jsonDecoder(avroSchema, json);
      avroPojos.add(specificReader.read(null, decoder).freeze());
    }
    ByteArrayOutputStream osBatch = new ByteArrayOutputStream();
    DatumWriter specificWriter = new SpecificDatumWriter<>(avroClass);
    try (DataFileWriter fileWriter =
             new DataFileWriter<>(specificWriter).create(avroSchema, osBatch)) {
      for (LogEventAvro avro : avroPojos) {
        fileWriter.append(avro);
        ByteArrayOutputStream osBinary = new ByteArrayOutputStream();
        Encoder encoder = EncoderFactory.get().directBinaryEncoder(osBinary, null);
        specificWriter.write(avro, encoder);
        avroBinaries.add(osBinary.toByteArray());
      }
    }
    final byte[] avroBatch = osBatch.toByteArray();
    for (GenericRecord generic : new DataFileStream<>(new ByteArrayInputStream(avroBatch),
                                                      new GenericDatumReader())) {
        avroGenerics.add(generic);
    }
    return avroBatch;
  }

  private void buildThrift() throws TException {
    TTestMemoryOutputTransport out = new TTestMemoryOutputTransport();
    for (byte[] json : THRIFT.jsonBytes) {
      LogEventThrift event = thriftSupplier.get();
      event.read(new TJSONProtocol(new TMemoryInputTransport(json)));
      thriftPojos.add(event);
      out.baos.reset();
      event.write(new TBinaryProtocol(out));
      thriftBinaries.add(out.baos.toByteArray());
      out.baos.reset();
      event.write(new TCompactProtocol(out));
      thriftCompacts.add(out.baos.toByteArray());
    }
  }

  private void buildProtobuf() throws IOException {
    JsonFormat jsonFormat = new JsonFormat();
    for (String json : PROTOBUF.jsonStrings) {
      AdExchangeProtobuf.LogEvent.Builder builder = protobufSupplier.get();
      jsonFormat.merge(json, ExtensionRegistry.getEmptyRegistry(), builder);
      AdExchangeProtobuf.LogEvent event = builder.build();
      protobufPojos.add(event);
      protobufBinaries.add(event.toByteArray());
    }
  }

  /* JSON */

  public Stream json() {
    return JSON.strings();
  }


  /* AVRO */

  public Stream avroSpecific() {
    return avroPojos.stream().sequential();
  }

  public Stream avroGeneric() {
    return avroGenerics.stream().sequential();
  }

  public Stream avroBinary() {
    return avroBinaries.stream().sequential();
  }

  public Stream avroJson() {
    return AVRO.strings();
  }

  public InputStream avroBatch() {
    return new ByteArrayInputStream(avroBatch);
  }


  /* THRIFT */

  public Stream thrift() {
    return thriftPojos.stream().sequential();
  }

  public Stream thriftBinary() {
    return thriftBinaries.stream().sequential();
  }

  public Stream thriftCompact() {
    return thriftCompacts.stream().sequential();
  }

  public Stream thriftJson() {
    return THRIFT.strings();
  }


  /* PROTOBUF */

  public Stream protobuf() {
    return protobufPojos.stream().sequential();
  }

  public Stream protobufBinary() {
    return protobufBinaries.stream().sequential();
  }

  public Stream protobufJson() {
    return PROTOBUF.strings();
  }


  /* */

  static public final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

  protected final List jsonObjects = json()
      .map(Unchecked.function(OBJECT_MAPPER::readTree))
      .collect(Collectors.toList());


  public JsonParser jsonNullsObjectParser() {
    String json = "{\"properties\":null,\"timestamp\":null,\"type\":null,\"uuid\":null,\"request\":null,\"response\":null}";
    return Unchecked.supplier(() -> AnoaTestSample.OBJECT_MAPPER.getFactory().createParser(json))
        .get();
  }

  public Stream jsonNodes() {
    return jsonObjects.stream().sequential();
  }

  public Stream jsonObjects() {
    return jsonNodes().map(node -> (ObjectNode) node);
  }

  public InputStream allAsInputStream(CheckedConsumer lambda) {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    Unchecked.consumer(lambda).accept(baos);
    return new ByteArrayInputStream(baos.toByteArray());
  }


  public void assertJsonNodes(Stream stream) {
    ListAssert.assertEquals(jsonObjects, stream.collect(Collectors.toList()));
  }

  public void assertJsonObjects(Stream stream) {
    ListAssert.assertEquals(jsonObjects, stream.collect(Collectors.toList()));
  }

  public InputStream jsonInputStream(int readFailureIndex) {
    return new TestInputStream(json().map(String::getBytes), readFailureIndex);
  }

  public JsonParser jsonParser(int readFailureIndex) {
    try {
      return OBJECT_MAPPER.getFactory().createParser(jsonInputStream(readFailureIndex));
    } catch (IOException e) {
      throw new UncheckedIOException(e);
    }
  }

  public void assertAvroGenerics(Stream stream) {
    ListAssert.assertEquals(avroGenerics,
                            stream.collect(Collectors.toList()));
  }

  public void assertAvroSpecifics(Stream stream) {
    ListAssert.assertEquals(avroPojos, stream.collect(Collectors.toList()));
  }

  public InputStream avroBinaryInputStream(int readFailureIndex) {
    return new TestInputStream(avroBinary(), readFailureIndex);
  }
  public InputStream avroJsonInputStream(int readFailureIndex) {
    return new TestInputStream(avroJson().map(String::getBytes), readFailureIndex);
  }

  public InputStream avroBatchInputStream(int readFailureIndex) {
    return new TestInputStream(avroBatch, readFailureIndex);
  }

  public void assertThriftObjects(Stream stream) {
    ListAssert.assertEquals(thriftPojos, stream.collect(Collectors.toList()));
  }

  public InputStream thriftBinaryInputStream(int readFailureIndex) {
    return new TestInputStream(thriftBinary(), readFailureIndex);
  }

  public InputStream thriftCompactInputStream(int readFailureIndex) {
    return new TestInputStream(thriftCompact(), readFailureIndex);
  }

  public InputStream thriftJsonInputStream(int readFailureIndex) {
    return new TestInputStream(thriftJson().map(String::getBytes), readFailureIndex);
  }

  public InputStream protoBinaryInputStream(int readFailureIndex) {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    protobuf().forEach(Unchecked.consumer(br -> br.writeDelimitedTo(baos)));
    return new TestInputStream(baos.toByteArray(), readFailureIndex);
  }

  public void assertProtobufObjects(Stream stream) {
    ListAssert.assertEquals(protobufPojos, stream.collect(Collectors.toList()));
  }
}