Examples with State org.apache.gobblin.configuration.State used on opensource projects

Example 16 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class ThrottleWriterTest method testGetFinalState.

public void testGetFinalState() throws IOException {
    PartitionedDataWriter writer = mock(PartitionedDataWriter.class);
    when(writer.getFinalState()).thenReturn(new State());
    int parallelism = 2;
    int qps = 4;
    DataWriter<Void> throttleWriter = setup(writer, parallelism, qps, ThrottleType.QPS);
    State state = ((FinalState) throttleWriter).getFinalState();
    verify(writer, times(1)).getFinalState();
    Assert.assertTrue(state.contains(ThrottleWriter.THROTTLED_TIME_KEY));
}

Example 17 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class ThrottleWriterTest method setup.

private DataWriter<Void> setup(DataWriter<Void> writer, int parallelism, int rate, ThrottleType type) throws IOException {
    State state = new State();
    state.appendToSetProp(ThrottleWriter.WRITER_LIMIT_RATE_LIMIT_KEY, Integer.toString(rate * parallelism));
    state.appendToSetProp(ThrottleWriter.WRITER_THROTTLE_TYPE_KEY, type.name());
    state.appendToSetProp(ConfigurationKeys.TASK_EXECUTOR_THREADPOOL_SIZE_KEY, Integer.toString(parallelism));
    state.appendToSetProp(ConfigurationKeys.SOURCE_MAX_NUMBER_OF_PARTITIONS, Integer.toString(parallelism));
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, state);
    return builder.build();
}

Example 18 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class SalesforceRestWriterTest method testBatchInsertDuplicate.

public void testBatchInsertDuplicate() throws IOException, URISyntaxException {
    final int recordSize = 25;
    final int batchSize = recordSize;
    State state = new State();
    state.appendToSetProp(CONF_PREFIX + BATCH_SIZE, Integer.toString(batchSize));
    state.appendToSetProp(CONF_PREFIX + BATCH_RESOURCE_PATH, "test");
    setup(SalesforceRestWriter.Operation.INSERT_ONLY_NOT_EXIST, state);
    CloseableHttpResponse response = mock(CloseableHttpResponse.class);
    StatusLine statusLine = mock(StatusLine.class);
    when(client.execute(any(HttpUriRequest.class))).thenReturn(response);
    when(response.getStatusLine()).thenReturn(statusLine);
    when(statusLine.getStatusCode()).thenReturn(200);
    HttpEntity entity = mock(HttpEntity.class);
    when(response.getEntity()).thenReturn(entity);
    JsonObject jsonResponse = new JsonObject();
    jsonResponse.addProperty("hasErrors", true);
    JsonArray resultJsonArr = new JsonArray();
    jsonResponse.add("results", resultJsonArr);
    JsonObject subResult1 = new JsonObject();
    subResult1.addProperty("statusCode", 400);
    JsonArray subResultArr = new JsonArray();
    JsonObject errJson = new JsonObject();
    errJson.addProperty("errorCode", SalesforceRestWriter.DUPLICATE_VALUE_ERR_CODE);
    subResultArr.add(errJson);
    subResult1.add("result", subResultArr);
    JsonObject subResult2 = new JsonObject();
    subResult2.addProperty("statusCode", 400);
    subResult2.add("result", subResultArr);
    resultJsonArr.add(subResult1);
    resultJsonArr.add(subResult2);
    when(entity.getContent()).thenReturn(new ByteArrayInputStream(jsonResponse.toString().getBytes()));
    RestEntry<JsonObject> restEntry = new RestEntry<JsonObject>("test", new JsonObject());
    writer = spy(writer);
    for (int i = 0; i < recordSize; i++) {
        writer.write(restEntry);
    }
    writer.commit();
    Assert.assertEquals(writer.recordsWritten(), recordSize);
    verify(writer, times(recordSize)).writeImpl(restEntry);
    verify(writer, times(recordSize)).onNewRecord(restEntry);
    double sendCount = ((double) recordSize) / ((double) batchSize);
    sendCount = Math.ceil(sendCount);
    verify(writer, times((int) sendCount)).sendRequest(any(HttpUriRequest.class));
    verify(client, times((int) sendCount)).execute(any(HttpUriRequest.class));
    verify(writer, times((int) sendCount)).waitForResponse(any(ListenableFuture.class));
    verify(writer, times((int) sendCount)).processResponse(any(CloseableHttpResponse.class));
    verify(writer, times(1)).flush();
    verify(writer, never()).onConnect(any(URI.class));
}

Example 19 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class SalesforceRestWriterTest method testBatchInsertSuccess.

public void testBatchInsertSuccess() throws IOException, URISyntaxException {
    final int recordSize = 113;
    final int batchSize = 25;
    State state = new State();
    state.appendToSetProp(CONF_PREFIX + BATCH_SIZE, Integer.toString(batchSize));
    state.appendToSetProp(CONF_PREFIX + BATCH_RESOURCE_PATH, "test");
    setup(SalesforceRestWriter.Operation.INSERT_ONLY_NOT_EXIST, state);
    CloseableHttpResponse response = mock(CloseableHttpResponse.class);
    StatusLine statusLine = mock(StatusLine.class);
    when(client.execute(any(HttpUriRequest.class))).thenReturn(response);
    when(response.getStatusLine()).thenReturn(statusLine);
    when(statusLine.getStatusCode()).thenReturn(200);
    HttpEntity entity = mock(HttpEntity.class);
    when(response.getEntity()).thenReturn(entity);
    JsonObject jsonResponse = new JsonObject();
    jsonResponse.addProperty("hasErrors", false);
    ByteArrayInputStream[] streams = new ByteArrayInputStream[recordSize];
    for (int i = 0; i < recordSize - 1; i++) {
        streams[i] = new ByteArrayInputStream(jsonResponse.toString().getBytes());
    }
    when(entity.getContent()).thenReturn(new ByteArrayInputStream(jsonResponse.toString().getBytes()), streams);
    RestEntry<JsonObject> restEntry = new RestEntry<JsonObject>("test", new JsonObject());
    writer = spy(writer);
    for (int i = 0; i < recordSize; i++) {
        writer.write(restEntry);
    }
    writer.commit();
    Assert.assertEquals(writer.recordsWritten(), recordSize);
    verify(writer, times(recordSize)).writeImpl(restEntry);
    verify(writer, times(recordSize)).onNewRecord(restEntry);
    double sendCount = ((double) recordSize) / ((double) batchSize);
    sendCount = Math.ceil(sendCount);
    verify(writer, times((int) sendCount)).sendRequest(any(HttpUriRequest.class));
    verify(client, times((int) sendCount)).execute(any(HttpUriRequest.class));
    verify(writer, times((int) sendCount)).waitForResponse(any(ListenableFuture.class));
    verify(writer, times((int) sendCount)).processResponse(any(CloseableHttpResponse.class));
    verify(writer, times(1)).flush();
    verify(writer, never()).onConnect(any(URI.class));
}

Example 20 with State

use of org.apache.gobblin.configuration.State in project incubator-gobblin by apache.

the class TimeBasedAvroWriterPartitionerTest method setUp.

@BeforeClass
public void setUp() throws IOException {
    File stagingDir = new File(STAGING_DIR);
    File outputDir = new File(OUTPUT_DIR);
    if (!stagingDir.exists()) {
        stagingDir.mkdirs();
    } else {
        FileUtils.deleteDirectory(stagingDir);
    }
    if (!outputDir.exists()) {
        outputDir.mkdirs();
    } else {
        FileUtils.deleteDirectory(outputDir);
    }
    this.schema = new Schema.Parser().parse(AVRO_SCHEMA);
    State properties = new State();
    properties.setProp(TimeBasedAvroWriterPartitioner.WRITER_PARTITION_COLUMNS, PARTITION_COLUMN_NAME);
    properties.setProp(ConfigurationKeys.WRITER_BUFFER_SIZE, ConfigurationKeys.DEFAULT_BUFFER_SIZE);
    properties.setProp(ConfigurationKeys.WRITER_FILE_SYSTEM_URI, ConfigurationKeys.LOCAL_FS_URI);
    properties.setProp(ConfigurationKeys.WRITER_STAGING_DIR, STAGING_DIR);
    properties.setProp(ConfigurationKeys.WRITER_OUTPUT_DIR, OUTPUT_DIR);
    properties.setProp(ConfigurationKeys.WRITER_FILE_PATH, BASE_FILE_PATH);
    properties.setProp(ConfigurationKeys.WRITER_FILE_NAME, FILE_NAME);
    properties.setProp(TimeBasedWriterPartitioner.WRITER_PARTITION_PATTERN, "yyyy/MM/dd");
    properties.setProp(ConfigurationKeys.WRITER_PARTITIONER_CLASS, TimeBasedAvroWriterPartitioner.class.getName());
    // Build a writer to write test records
    DataWriterBuilder<Schema, GenericRecord> builder = new AvroDataWriterBuilder().writeTo(Destination.of(Destination.DestinationType.HDFS, properties)).writeInFormat(WriterOutputFormat.AVRO).withWriterId(WRITER_ID).withSchema(this.schema).withBranches(1).forBranch(0);
    this.writer = new PartitionedDataWriter<Schema, GenericRecord>(builder, properties);
}