Unit Testing DoFns & Transforms

1. Introduction

2. Why This Concept Exists

When a pipeline fails in production due to a parsing bug or edge-case null value:

• Running an entire 20-step pipeline just to test a single 5-line string parsing function is extremely slow.
• Isolating the specific DoFn class allows software engineers to execute unit test assertions in under 100 milliseconds on their local laptops using standard Python test frameworks like unittest or pytest.

3. Key Terminology

• Unit Test: A fast, isolated test evaluating a single function or class method without external network dependencies.
• Mocking: Replacing external database or web service calls with synthetic fake responses during unit test execution.
• DoFn.process Unit Test: Verifying that a DoFn emits exact expected outputs when given specific input elements.

4. How It Works

1. Instantiate Target DoFn: Create an instance of your custom DoFn class in a test method.
2. Define Edge-Case Inputs: Create lists of inputs including normal records, empty strings, null values, and negative numbers.
3. Run inside TestPipeline: Pass synthetic inputs through beam.Create and beam.ParDo(MyDoFn()).
4. Assert Output Matching: Use assert_that(results, equal_to(expected)) to verify that outputs match expected output records.

5. Visual Diagram

6. Code Example

pythonCopy
import unittest
import re
import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that, equal_to

# DoFn under test
class CleanPhoneNumberDoFn(beam.DoFn):
    def process(self, element):
        # Strip all non-digit characters
        digits_only = re.sub(r"\D", "", element)
        # Yield only if phone number has exactly 10 digits
        if len(digits_only) == 10:
            yield f"+1{digits_only}"

class TestCleanPhoneNumberDoFn(unittest.TestCase):

    def test_valid_and_invalid_phone_numbers(self):
        inputs = [
            "(555) 123-4567", # Valid 10 digits
            "555.987.6543",   # Valid 10 digits
            "123-45",         # Invalid (too short)
            "invalid-string"   # Invalid (no digits)
        ]
        expected_outputs = [
            "+15551234567",
            "+15559876543"
        ]

        with TestPipeline() as p:
            results = (
                p
                | "CreateTestPhoneNumbers" >> beam.Create(inputs)
                | "CleanPhones" >> beam.ParDo(CleanPhoneNumberDoFn())
            )

            assert_that(results, equal_to(expected_outputs))

if __name__ == "__main__":
    unittest.main()

7. Code Explanation

• CleanPhoneNumberDoFn strips non-digits using regex and yields normalized +1XXXXXXXXXX strings.
• test_valid_and_invalid_phone_numbers tests both valid formatted numbers and invalid short/alphabetic strings.
• assert_that(results, equal_to(expected_outputs)) verifies that invalid inputs are safely dropped while valid numbers are correctly formatted.

8. Real Production Example

In a customer fraud system:

• Engineers write isolated unit tests for FraudFilterDoFn to test edge cases: negative transaction amounts, missing currency codes, and unexpected string symbols.
• Because tests execute locally in milliseconds, developers run pytest continuously while editing code (TDD workflow).

9. Common Mistakes

• Making live network calls inside unit tests: Invoking real external REST APIs or database queries inside a unit test causes slow, fragile test runs that fail when offline. Mock all external dependencies.
• Omitting edge cases: Testing only "happy path" inputs and forgetting null/empty inputs.

10. Best Practices

• Write unit tests for every custom DoFn class before combining them into composite pipelines.
• Mock external network calls using Python's unittest.mock module.

11. Summary

• Unit testing validates isolated DoFn business logic in under 100 milliseconds.
• Uses TestPipeline and beam.Create to feed synthetic test records.
• Ensures code is robust against edge cases before cloud deployment.

12. Interactive Challenges

13. Related Content

• Testing Beam Pipelines
• Integration Testing