Testing I/O Transforms in Apache Beam

Examples and design patterns for testing Apache Beam I/O transforms

Adapt for:

Java SDK

Python SDK

Note: This guide is still in progress. There is an open issue to finish the guide: BEAM-1025.

Introduction

This document explains the set of tests that the Beam community recommends based on our past experience writing I/O transforms. If you wish to contribute your I/O transform to the Beam community, we’ll ask you to implement these tests.

While it is standard to write unit tests and integration tests, there are many possible definitions. Our definitions are:

Data store used: an in-memory version of the data store (if available), otherwise you’ll need to write a fake

Data set size: tiny (10s to 100s of rows)

Integration Tests:

Goal: catch problems that occur when interacting with real versions of the runners/data store

Data store used: an actual instance, pre-configured before the test

Data set size: small to medium (1000 rows to 10s of GBs)

A note on performance benchmarking

We do not advocate writing a separate test specifically for performance benchmarking. Instead, we recommend setting up integration tests that can accept the necessary parameters to cover many different testing scenarios.

For example, if integration tests are written according to the guidelines below, the integration tests can be run on different runners (either local or in a cluster configuration) and against a data store that is a small instance with a small data set, or a large production-ready cluster with larger data set. This can provide coverage for a variety of scenarios - one of them is performance benchmarking.

Test Balance - Unit vs Integration

It’s easy to cover a large amount of code with an integration test, but it is then hard to find a cause for test failures and the test is flakier.

However, there is a valuable set of bugs found by tests that exercise multiple workers reading/writing to data store instances that have multiple nodes (eg, read replicas, etc.). Those scenarios are hard to find with unit tests and we find they commonly cause bugs in I/O transforms.

Our test strategy is a balance of those 2 contradictory needs. We recommend doing as much testing as possible in unit tests, and writing a single, small integration test that can be run in various configurations.

Examples

Java:

BigtableIO’s testing implementation is considered the best example of current best practices for unit testing Sources

JdbcIO has the current best practice examples for writing integration tests.

Unit Tests

Goals

Validate the correctness of the code in your I/O transform.

Validate that the I/O transform works correctly when used in concert with reference implementations of the data store it connects with (where “reference implementation” means a fake or in-memory version).

Be able to run quickly and need only one machine, with a reasonably small memory/disk footprint and no non-local network access (preferably none at all). Aim for tests than run within several seconds - anything above 20 seconds should be discussed with the beam dev mailing list.

Validate that the I/O transform can handle network failures.

Non-goals

Test problems in the external data store - this can lead to extremely complicated tests.

Implementing unit tests

A general guide to writing Unit Tests for all transforms can be found in the PTransform Style Guide. We have expanded on a few important points below.

If you are using the Source API, make sure to exhaustively unit-test your code. A minor implementation error can lead to data corruption or data loss (such as skipping or duplicating records) that can be hard for your users to detect. Also look into using SourceTestUtilssource_test_utils - it is a key piece of testing Source implementations.

If you are not using the Source API, you can use TestPipeline with PAssertassert_that to help with your testing.

If you are implementing write, you can use TestPipeline to write test data and then read and verify it using a non-Beam client.

Use fakes

Instead of using mocks in your unit tests (pre-programming exact responses to each call for each test), use fakes. The preferred way to use fakes for I/O transform testing is to use a pre-existing in-memory/embeddable version of the service you’re testing, but if one does not exist consider implementing your own. Fakes have proven to be the right mix of “you can get the conditions for testing you need” and “you don’t have to write a million exacting mock function calls”.

Network failure

To help with testing and separation of concerns, code that interacts across a network should be handled in a separate class from your I/O transform. The suggested design pattern is that your I/O transform throws exceptions once it determines that a read or write is no longer possible.

This allows the I/O transform’s unit tests to act as if they have a perfect network connection, and they do not need to retry/otherwise handle network connection problems.

Batching

If your I/O transform allows batching of reads/writes, you must force the batching to occur in your test. Having configurable batch size options on your I/O transform allows that to happen easily. These must be marked as test only.

I/O Transform Integration Tests

We do not currently have examples of Python I/O integration tests or integration tests for unbounded or eventually consistent data stores. We would welcome contributions in these areas - please contact the Beam dev@ mailing list for more information.

Goals

Allow end to end testing of interactions between data stores, I/O transforms, and runners, simulating real world conditions.

Allow both small scale and large scale testing.

Self contained: require the least possible initial setup or existing outside state, besides the existence of a data store that the test can modify.

Anyone can run the same set of I/O transform integration tests that Beam runs on its continuous integration servers.

Integration tests, data stores, and Kubernetes

In order to test I/O transforms in real world conditions, you must connect to a data store instance.

The Beam community hosts the data stores used for integration tests in Kubernetes. In order for an integration test to be run in Beam’s continuous integration environment, it must have Kubernetes scripts that set up an instance of the data store.

However, when working locally, there is no requirement to use Kubernetes. All of the test infrastructure allows you to pass in connection info, so developers can use their preferred hosting infrastructure for local development.

Running integration tests

The high level steps for running an integration test are:

Set up the data store corresponding to the test being run

Run the test, passing it connection info from the just created data store

Clean up the data store

Since setting up data stores and running the tests involves a number of steps, and we wish to time these tests when running performance benchmarks, we use PerfKit Benchmarker to manage the process end to end. With a single command, you can go from an empty Kubernetes cluster to a running integration test.

However, PerfKit Benchmarker is not required for running integration tests. Therefore, we have listed the steps for both using PerfKit Benchmarker, and manually running the tests below.

Using PerfKit Benchmarker

Have a running Kubernetes cluster you can connect to locally using kubectl

You won’t need to invoke PerfKit Benchmarker directly. Run mvn verify in the directory of the I/O module you’d like to test, with the parameter io-it-suite when running in jenkins CI or with a kubernetes cluster on the same network or io-it-suite-local when running on a local dev box accessing a kubernetes cluster on a remote network.

Run the test using the instructions in the class (e.g. see the instructions in JdbcIOIT.java)

Tell Kubernetes to delete the resources specified in the Kubernetes scripts:

JDBC: kubectl delete -f .test-infra/kubernetes/postgres/postgres.yml

Elasticsearch: bash .test-infra/kubernetes/elasticsearch/teardown.sh

Implementing Integration Tests

There are three components necessary to implement an integration test:

Test code: the code that does the actual testing: interacting with the I/O transform, reading and writing data, and verifying the data.

Kubernetes scripts: a Kubernetes script that sets up the data store that will be used by the test code.

Integrate with PerfKit Benchmarker using io-it-suite: this allows users to easily invoke PerfKit Benchmarker, creating the Kubernetes resources and running the test code.

These three pieces are discussed in detail below.

Test Code

These are the conventions used by integration testing code:

Your test should use pipeline options to receive connection information.

For Java, there is a shared pipeline options object in the io/common directory. This means that if there are two tests for the same data store (e.g. for Elasticsearch and the HadoopInputFormatIO tests), those tests share the same pipeline options.

Generate test data programmatically and parameterize the amount of data used for testing.

For Java, CountingInput + TestRow can be combined to generate deterministic test data at any scale.

Use a write then read style for your tests.

In a single Test, run a pipeline to do a write using your I/O transform, then run another pipeline to do a read using your I/O transform.

The only verification of the data should be the result from the read. Don’t validate the data written to the database in any other way.

Validate the actual contents of all rows in an efficient manner. An easy way to do this is by taking a hash of the rows and combining them. HashingFn can help make this simple, and TestRow has pre-computed hashes.

For easy debugging, use PAssert’s containsInAnyOrder to validate the contents of a subset of all rows.

Tests should assume they may be run multiple times and/or simultaneously on the same database instance.

Clean up test data: do this in an @AfterClass to ensure it runs.

Use unique table names per run (timestamps are an easy way to do this) and per-method where appropriate.

If you would like help with this or have other questions, contact the Beam dev@ mailing list and the community may be able to assist you.

Guidelines for creating a Beam data store Kubernetes script:

You must only provide access to the data store instance via a NodePort service.

This is a requirement for security, since it means that only the local network has access to the data store. This is particularly important since many data stores don’t have security on by default, and even if they do, their passwords will be checked in to our public Github repo.

You should define two Kubernetes scripts.

This is the best known way to implement item #1.

The first script will contain the main datastore instance script (StatefulSet) plus a NodePort service exposing the data store. This will be the script run by the Beam Jenkins continuous integration server.

The second script will define a LoadBalancer service, used for local development if the Kubernetes cluster is on another network. This file’s name is usually suffixed with ‘-for-local-dev’.

You must ensure that pods are recreated after crashes.

If you use a pod directly, it will not be recreated if the pod crashes or something causes the cluster to move the container for your pod.

In most cases, you’ll want to use StatefulSet as it supports persistent disks that last between restarts, and having a stable network identifier associated with the pod using a particular persistent disk. Deployment and ReplicaSet are also possibly useful, but likely in fewer scenarios since they do not have those features.

You should create separate scripts for small and large instances of your data store.

You must use a Docker image from a trusted source and pin the version of the Docker image.

You should prefer images in this order:

An image provided by the creator of the data source/sink (if they officially maintain it). For Apache projects, this would be the official Apache repository.

Official Docker images, because they have security fixes and guaranteed maintenance.

Non-official Docker images, or images from other providers that have good maintainers (e.g. quay.io).

Integrate with PerfKit Benchmarker

To allow developers to easily invoke your I/O integration test, you must perform these two steps. The follow sections describe each step in more detail.

Create a PerfKit Benchmarker benchmark configuration file for the data store. Each pipeline option needed by the integration test should have a configuration entry.

Modify the per-I/O Maven pom configuration so that PerfKit Benchmarker can be invoked from Maven.

The goal is that a checked in config has defaults such that other developers can run the test without changing the configuration.

Defining the benchmark configuration file

The benchmark configuration file is a yaml file that defines the set of pipeline options for a specific data store. Some of these pipeline options are static - they are known ahead of time, before the data store is created (e.g. username/password). Others options are dynamic - they are only known once the data store is created (or after we query the Kubernetes cluster for current status).

All known cases of dynamic pipeline options are for extracting the IP address that the test needs to connect to. For I/O integration tests, we must allow users to specify:

The type of the IP address to get (load balancer/node address)

The pipeline option to pass that IP address to

How to find the Kubernetes resource with that value (ie. what load balancer service name? what node selector?)

The style of dynamic pipeline options used here should support a variety of other types of values derived from Kubernetes, but we do not have specific examples.

The dynamic pipeline options are:

Type name

Meaning

Selector field name

Selector field value

NodePortIp

We will be using the IP address of a k8s NodePort service, the value will be an IP address of a Pod

podLabel

A kubernetes label selector for a pod whose IP address can be used to connect to

LoadBalancerIp

We will be using the IP address of a k8s LoadBalancer, the value will be an IP address of the load balancer

The set of mvn pipeline options that PerfKit Benchmarker will determine at runtime.

dynamic_pipeline_options.name

The name of the parameter to be passed to mvn's invocation of the I/O integration test.

dynamic_pipeline_options.type

The method of determining the value of the pipeline options.

dynamic_pipeline_options - other attributes

These vary depending on the type of the dynamic pipeline option - see the table of dynamic pipeline options for a description.

Per-I/O mvn pom configuration

Each I/O is responsible for adding a section to its pom with a profile that invokes PerfKit Benchmarker with the proper parameters during the verify phase. Below are the set of PerfKit Benchmarker parameters and how to configure them.

The JdbcIO pom has an example of how to put these options together into a profile and invoke Python+PerfKit Benchmarker with them.

PerfKit Benchmarker Parameter

Description

Example value

benchmarks

Defines the PerfKit Benchmarker benchmark to run. This is same for all I/O integration tests.

beam_integration_benchmark

beam_location

The location where PerfKit Benchmarker can find the Beam repository.

${beamRootProjectDir} - this is a variable you'll need to define for each maven pom. See example pom for an example.

beam_prebuilt

Whether or not to rebuild the Beam repository before invoking the I/O integration test command.

true

beam_sdk

Whether PerfKit Benchmarker will run the Beam SDK for Java or Python.

java

beam_runner_profile

Optional command line parameter used to override the runner, allowing us to use the direct runner.

Always use the predefined variable instead of specifying this parameter ${pkbBeamRunnerProfile}

beam_runner_option

Optional command line parameter used to override the runner, allowing us to use the direct runner.

Always use the predefined variable instead of specifying this parameter ${pkbBeamRunnerOption}

beam_it_module

The path to the pom that contains the test (needed for invoking the test with PerfKit Benchmarker).

sdks/java/io/jdbc

beam_it_class

The test to run.

org.apache.beam.sdk.io.jdbc.JdbcIOIT

beam_it_options

Pipeline options for the beam job - meant to be a way to pass pipeline options the user specifies on the commandline when invoking io-it-suite

Always use ${integrationTestPipelineOptions}, which allows the user to pass in parameters.

There is also a set of Maven properties which are useful when invoking PerfKit Benchmarker. These properties are configured in the I/O parent pom, and some are only available when the io-it-suite profile is active in Maven.

Small Scale and Large Scale Integration Tests

Apache Beam expects that it can run integration tests in multiple configurations:

Small scale

Execute on a single worker on the runner (it should be possible but is not required).

The data store should be configured to use a single node.

The dataset can be very small (1000 rows).

Large scale

Execute on multiple workers on the runner.

The datastore should be configured to use multiple nodes.

The data set used in this case is larger (10s of GBs).

You can do this by:

Creating two Kubernetes scripts: one for a small instance of the data store, and one for a large instance.

Having your test take a pipeline option that decides whether to generate a small or large amount of test data (where small and large are sizes appropriate to your data store)

Apache Beam, Apache, Beam, the Beam logo, and the Apache feather logo are
either registered trademarks or trademarks of The Apache Software
Foundation. All other products or name brands are trademarks of their
respective holders, including The Apache Software Foundation.