How to Unit Test CloudBigTableIO with Dataflow Pipeline?

user_4487887 0
Viewed 144

I have a Dataflow pipeline (Apache Beam) that reads from Kafka, converts to Put (Mutation) and writes to BigTable using the CloudBigtableIO.writeToTable().

I wrote a unit test to validate this conversion and test my DoFn class KafkaToBigtableRow. I tried to compare the result using PAssert, like:

    @Test
    public void kafkaToBigtable() {

        final List<KV<Long, String>> kafkaMessages = ImmutableList.of(
                prepareKV(1L, "123", "guid123", "123abc"),
                prepareKV(2L, "456", "guid546", "456abc")
        );

        final PCollection<Mutation> output = testPipeline
                .apply(Create.of(kafkaMessages))
                .apply(ParDo.of(new KafkaToBigtableRow()));

        final Put[] mutations = {
                new Put(Bytes.toBytes(prepareRowKey("123", "123abc"))).addColumn(COLUMN_FAMILY, GUID_KEY.getBytes(), "guid123".getBytes()),
                new Put(Bytes.toBytes(prepareRowKey("456", "456abc"))).addColumn(COLUMN_FAMILY, GUID_KEY.getBytes(), "guid456".getBytes())
        };

        PAssert.that(output).containsInAnyOrder(mutations);

        testPipeline.run();
    }

But I get the assertion error:

java.lang.AssertionError: ParDo(KafkaToBigtableRow)/ParMultiDo(KafkaToBigtableRow).output: 
Expected: iterable with items [<{"totalColumns":1,"row":"123#123abc#20201203","families":{"b":[{"qualifier":"guid","vlen":7,"tag":[],"timestamp":"9223372036854775807"}]}}>, <{"totalColumns":1,"row":"456#456abc#20201203","families":{"b":[{"qualifier":"guid","vlen":7,"tag":[],"timestamp":"9223372036854775807"}]}}>] in any order
     but: not matched: <{"totalColumns":1,"row":"456#456abc#20201203","families":{"b":[{"qualifier":"guid","vlen":7,"tag":[],"timestamp":"9223372036854775807"}]}}>

I compared the strings and they are equals. What I could find is that PAssert uses the equals method to compare the Objects and the class org.apache.hadoop.hbase.client.Put does not implement equals, so it uses the default from Object and returns false.

Is there a way to workaround this issue? Or is there another way to test this mutation?

I know there's the option to to am integration test and use a BigTable emulator, but I want to avoid that because of portability.

Thanks!

apache-beam dataflow bigtable pipeline unit-testing
1 Answers

I was able to fix the problem using PAssert satisfies() with a custom SerializableFuntion.

@Test
public void kafkaToBigtable() {

    final List<KV<Long, String>> kafkaMessages = ImmutableList.of(
            prepareKV(1L, "123", "guid123", "123abc"),
            prepareKV(2L, "456", "guid546", "456abc")
    );

    final PCollection<Mutation> output = testPipeline
            .apply(Create.of(kafkaMessages))
            .apply(ParDo.of(new KafkaToBigtableRow()));

    PAssert.that(output).satisfies(VerifyAnySample.validateMutation());

    testPipeline.run();
}

private abstract static class VerifyAnySample implements SerializableFunction<Iterable<Mutation>, Void> {

    private static SerializableFunction<Iterable<Mutation>, Void> validateMutation() {
        return input -> {
            assertValues(input.iterator());
            return null;
        };
    }

    private static void assertValues(Iterator<Mutation> iterator) {
        final Put mutation = (Put) iterator.next();
        final List<Cell> cellList = mutation.get(COLUMN_FAMILY, COLUMN_QUALIFIER);

        cellList.forEach(cell -> {

            final String rowKey = Bytes.toString(cell.getRowArray()).substring(cell.getRowOffset(), cell.getRowOffset() + cell.getRowLength());
            final String value = Bytes.toString(cell.getValueArray()).substring(cell.getValueOffset());

            // Logic to validate rowkey and value
        });

        if (iterator.hasNext()) {
            assertValues(iterator);
        }
    }
}

Maybe this is not the best way to do it, but it works.

user_44878870
Related