com.streamsets.pipeline.lib.parser.protobuf.TestProtobufDataParser.java Source code

Java tutorial

Introduction

Here is the source code for com.streamsets.pipeline.lib.parser.protobuf.TestProtobufDataParser.java

Source

/**
 * Copyright 2015 StreamSets Inc.
 *
 * Licensed under 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 com.streamsets.pipeline.lib.parser.protobuf;

import com.google.common.io.Resources;
import com.streamsets.pipeline.api.Field;
import com.streamsets.pipeline.api.OnRecordError;
import com.streamsets.pipeline.api.Record;
import com.streamsets.pipeline.api.Stage;
import com.streamsets.pipeline.lib.parser.DataParser;
import com.streamsets.pipeline.lib.parser.DataParserException;
import com.streamsets.pipeline.lib.parser.DataParserFactory;
import com.streamsets.pipeline.lib.parser.DataParserFactoryBuilder;
import com.streamsets.pipeline.lib.parser.DataParserFormat;
import com.streamsets.pipeline.lib.parser.Errors;
import com.streamsets.pipeline.lib.util.PersonProto;
import com.streamsets.pipeline.lib.util.ProtobufConstants;
import com.streamsets.pipeline.lib.util.ProtobufTestUtil;
import com.streamsets.pipeline.sdk.ContextInfoCreator;
import org.apache.commons.io.FileUtils;
import org.junit.Assert;
import org.junit.Test;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

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

public class TestProtobufDataParser {

    @Test
    public void testProtobufDataParser() throws IOException, DataParserException {
        DataParser dataParser = getDataParser("0", "Employee.desc", "util.Employee");
        List<Record> records = new ArrayList<>();
        Record r = dataParser.parse();
        // Engineer object in the stream is 59 bytes long
        assertEquals("138", dataParser.getOffset());
        while (r != null) {
            records.add(r);
            r = dataParser.parse();
        }
        assertTrue("-1".equals(dataParser.getOffset()));
        assertEquals(10, records.size());
        ProtobufTestUtil.compareProtoRecords(records, 0);
    }

    @Test
    public void testProtobuf3Delimited() throws Exception {
        DataParser dataParser = getDataParserFactory("TestRecordProtobuf3.desc", "TestRecord", true)
                .getParser("TestRecord", new ByteArrayInputStream(FileUtils
                        .readFileToByteArray(new File(Resources.getResource("TestProtobuf3.ser").getPath()))), "0");

        Record record = dataParser.parse();
        verifyProtobuf3TestRecord(record);
    }

    @Test
    public void testProtobuf3NonDelimited() throws Exception {
        DataParser dataParser = getDataParserFactory(
                "TestRecordProtobuf3.desc", "TestRecord", false)
                        .getParser("TestRecord",
                                new ByteArrayInputStream(FileUtils.readFileToByteArray(
                                        new File(Resources.getResource("TestProtobuf3NoDelimiter.ser").getPath()))),
                                "0");

        Record record = dataParser.parse();
        verifyProtobuf3TestRecord(record);
    }

    private void verifyProtobuf3TestRecord(Record record) {
        assertTrue(record.has("/first_name"));
        assertTrue(record.has("/full_name"));
        assertTrue(record.has("/test_map"));
        assertTrue(record.has("/samples"));

        assertEquals("Adam", record.get("/first_name").getValueAsString());
        assertEquals("", record.get("/full_name").getValue());
        List<Field> samples = record.get("/samples").getValueAsList();
        assertEquals(2, samples.size());
        assertEquals(1, samples.get(0).getValueAsInteger());
        assertEquals(2, samples.get(1).getValueAsInteger());
        Map<String, Field> testMap = record.get("/test_map").getValueAsMap();
        assertEquals(2, testMap.size());
        assertTrue(testMap.containsKey("hello"));
        assertTrue(testMap.containsKey("bye"));
        assertEquals("world", testMap.get("hello").getValueAsString());
        assertEquals("earth", testMap.get("bye").getValueAsString());
    }

    @Test
    public void testProtobufDataParserWithOffset() throws IOException, DataParserException {
        // Engineer object in the stream is 59 bytes long
        // Executive object in the stream is 54 bytes long
        // Skip first 5 objects => start from position 285 [59*3 54*2]
        DataParser dataParser = getDataParser("748", "Employee.desc", "util.Employee");
        List<Record> records = new ArrayList<>();
        assertEquals("748", dataParser.getOffset());
        Record r = dataParser.parse();
        assertEquals("915", dataParser.getOffset());
        while (r != null) {
            records.add(r);
            r = dataParser.parse();
        }
        assertTrue("-1".equals(dataParser.getOffset()));
        assertEquals(5, records.size());
        ProtobufTestUtil.compareProtoRecords(records, 5);
    }

    @Test
    public void testMissingRequiredField() throws Exception {
        ByteArrayOutputStream bOut = new ByteArrayOutputStream();
        PersonProto.Person john = PersonProto.Person.newBuilder().setId(1).setName("John Doe")
                .addEmail("jdoe" + "@example.com")
                .addPhone(PersonProto.Person.PhoneNumber.newBuilder().setNumber("7568345")).build();
        john.writeDelimitedTo(bOut);
        bOut.flush();

        DataParserFactory factory = getDataParserFactory("test1.desc", "util.Person");
        DataParser parser = factory.getParser("Person", new ByteArrayInputStream(bOut.toByteArray()), "0");

        try {
            parser.parse();
            Assert.fail("DataParserException expected as a required field is missing");
        } catch (DataParserException e) {
            assertEquals(Errors.DATA_PARSER_02, e.getErrorCode());
        }
    }

    public DataParser getDataParser(String offset, String protoFile, String messageType)
            throws IOException, DataParserException {
        return getDataParserFactory(protoFile, messageType).getParser("Person",
                new ByteArrayInputStream(ProtobufTestUtil.getProtoBufData()), offset);
    }

    private Stage.Context getContext() {
        return ContextInfoCreator.createSourceContext("i", false, OnRecordError.TO_ERROR,
                Collections.<String>emptyList());
    }

    public DataParserFactory getDataParserFactory(String protoFile, String messageType) {
        return getDataParserFactory(protoFile, messageType, true);
    }

    public DataParserFactory getDataParserFactory(String protoFile, String messageType, boolean isDelimited) {
        DataParserFactoryBuilder dataParserFactoryBuilder = new DataParserFactoryBuilder(getContext(),
                DataParserFormat.PROTOBUF);

        return dataParserFactoryBuilder
                .setConfig(ProtobufConstants.PROTO_DESCRIPTOR_FILE_KEY, Resources.getResource(protoFile).getPath())
                .setConfig(ProtobufConstants.MESSAGE_TYPE_KEY, messageType)
                .setConfig(ProtobufConstants.DELIMITED_KEY, isDelimited).setOverRunLimit(1000)
                .setMaxDataLen(Integer.MAX_VALUE).build();
    }
}