Examples with FakeBigQueryServices org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices used on opensource projects

Example 11 with FakeBigQueryServices

use of org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices in project beam by apache.

the class BigQueryIOStorageReadTest method testTableSourceEstimatedSize_WithDefaultProject.

@Test
public void testTableSourceEstimatedSize_WithDefaultProject() throws Exception {
    fakeDatasetService.createDataset("project-id", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("project-id:dataset.table");
    Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
    fakeDatasetService.createTable(table);
    BigQueryStorageTableSource<TableRow> tableSource = BigQueryStorageTableSource.create(ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")), null, null, new TableRowParser(), TableRowJsonCoder.of(), new FakeBigQueryServices().withDatasetService(fakeDatasetService));
    assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}

Example 12 with FakeBigQueryServices

use of org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices in project beam by apache.

the class BigQueryIOStorageReadTest method testReadFromStreamSource.

@Test
public void testReadFromStreamSource() throws Exception {
    ReadSession readSession = ReadSession.newBuilder().setName("readSession").setAvroSchema(AvroSchema.newBuilder().setSchema(AVRO_SCHEMA_STRING)).build();
    ReadRowsRequest expectedRequest = ReadRowsRequest.newBuilder().setReadStream("readStream").build();
    List<GenericRecord> records = Lists.newArrayList(createRecord("A", 1, AVRO_SCHEMA), createRecord("B", 2, AVRO_SCHEMA), createRecord("C", 3, AVRO_SCHEMA));
    List<ReadRowsResponse> responses = Lists.newArrayList(createResponse(AVRO_SCHEMA, records.subList(0, 2), 0.0, 0.50), createResponse(AVRO_SCHEMA, records.subList(2, 3), 0.5, 0.75));
    StorageClient fakeStorageClient = mock(StorageClient.class);
    when(fakeStorageClient.readRows(expectedRequest, "")).thenReturn(new FakeBigQueryServerStream<>(responses));
    BigQueryStorageStreamSource<TableRow> streamSource = BigQueryStorageStreamSource.create(readSession, ReadStream.newBuilder().setName("readStream").build(), TABLE_SCHEMA, new TableRowParser(), TableRowJsonCoder.of(), new FakeBigQueryServices().withStorageClient(fakeStorageClient));
    List<TableRow> rows = new ArrayList<>();
    BoundedReader<TableRow> reader = streamSource.createReader(options);
    for (boolean hasNext = reader.start(); hasNext; hasNext = reader.advance()) {
        rows.add(reader.getCurrent());
    }
    System.out.println("Rows: " + rows);
    assertEquals(3, rows.size());
}

Example 13 with FakeBigQueryServices

use of org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices in project beam by apache.

the class BigQueryIOStorageReadTest method testTableSourceEstimatedSize_WithBigQueryProject.

@Test
@ProjectOverride
public void testTableSourceEstimatedSize_WithBigQueryProject() throws Exception {
    fakeDatasetService.createDataset("bigquery-project-id", "dataset", "", "", null);
    TableReference tableRef = BigQueryHelpers.parseTableSpec("bigquery-project-id:dataset.table");
    Table table = new Table().setTableReference(tableRef).setNumBytes(100L);
    fakeDatasetService.createTable(table);
    BigQueryStorageTableSource<TableRow> tableSource = BigQueryStorageTableSource.create(ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("dataset.table")), null, null, new TableRowParser(), TableRowJsonCoder.of(), new FakeBigQueryServices().withDatasetService(fakeDatasetService));
    assertEquals(100, tableSource.getEstimatedSizeBytes(options));
}

Example 14 with FakeBigQueryServices

use of org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices in project beam by apache.

the class BigQueryIOStorageReadTest method testTableSourceCreateReader.

@Test
public void testTableSourceCreateReader() throws Exception {
    BigQueryStorageTableSource<TableRow> tableSource = BigQueryStorageTableSource.create(ValueProvider.StaticValueProvider.of(BigQueryHelpers.parseTableSpec("foo.com:project:dataset.table")), null, null, new TableRowParser(), TableRowJsonCoder.of(), new FakeBigQueryServices().withDatasetService(fakeDatasetService));
    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("BigQuery storage source must be split before reading");
    tableSource.createReader(options);
}

Example 15 with FakeBigQueryServices

use of org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices in project beam by apache.

the class BigQueryIOStorageReadTest method testStreamSourceSplit.

@Test
public void testStreamSourceSplit() throws Exception {
    BigQueryStorageStreamSource<TableRow> streamSource = BigQueryStorageStreamSource.create(ReadSession.getDefaultInstance(), ReadStream.getDefaultInstance(), TABLE_SCHEMA, new TableRowParser(), TableRowJsonCoder.of(), new FakeBigQueryServices());
    assertThat(streamSource.split(0, options), containsInAnyOrder(streamSource));
}