schrodinger.seam.testing.util module

A module that wraps apache_beam.testing.util while modifying some of its functions. with quality-of-life improvements.

schrodinger.seam.testing.util.assert_pcoll_size(pcoll: apache_beam.pvalue.PCollection, expected_size: int, rel_tolerance: float = 0.0, abs_tolerance: int = 0)

Asserts that the given PCollection has the expected size, with optional relative and absolute tolerances.

Example usages:

>>> with beam.Pipeline() as p:
...     elements = p | beam.Create([1, 2, 3])
...     assert_pcoll_size(elements, 3)
>>> with beam.Pipeline() as p:
...     elements = p | beam.Create(range(99))
...     assert_pcoll_size(elements, 100, rel_tolerance=0.01)

Args:

pcoll: The PCollection to check the size of. expected_size: The expected size of the PCollection. rel_tolerance: Optional, the relative tolerance allowed for the size of the PCollection. abs_tolerance: Optional, the absolute tolerance allowed for the size of the PCollection.