Search in sources :

Example 1 with FormatCountsFn

use of org.apache.beam.examples.cookbook.BigQueryTornadoes.FormatCountsFn in project beam by apache.

the class BigQueryTornadoesTest method testFormatCounts.

@Test
@SuppressWarnings({ "rawtypes", "unchecked" })
public void testFormatCounts() throws Exception {
    DoFnTester<KV<Integer, Long>, TableRow> formatCountsFn = DoFnTester.of(new FormatCountsFn());
    KV[] empty = {};
    List<TableRow> results = formatCountsFn.processBundle(empty);
    Assert.assertTrue(results.size() == 0);
    KV[] input = { KV.of(3, 0L), KV.of(4, Long.MAX_VALUE), KV.of(5, Long.MIN_VALUE) };
    results = formatCountsFn.processBundle(input);
    Assert.assertEquals(results.size(), 3);
    Assert.assertEquals(results.get(0).get("month"), 3);
    Assert.assertEquals(results.get(0).get("tornado_count"), 0L);
    Assert.assertEquals(results.get(1).get("month"), 4);
    Assert.assertEquals(results.get(1).get("tornado_count"), Long.MAX_VALUE);
    Assert.assertEquals(results.get(2).get("month"), 5);
    Assert.assertEquals(results.get(2).get("tornado_count"), Long.MIN_VALUE);
}
Also used : FormatCountsFn(org.apache.beam.examples.cookbook.BigQueryTornadoes.FormatCountsFn) TableRow(com.google.api.services.bigquery.model.TableRow) KV(org.apache.beam.sdk.values.KV) Test(org.junit.Test)

Aggregations

TableRow (com.google.api.services.bigquery.model.TableRow)1 FormatCountsFn (org.apache.beam.examples.cookbook.BigQueryTornadoes.FormatCountsFn)1 KV (org.apache.beam.sdk.values.KV)1 Test (org.junit.Test)1