Unit Testing Custom Code Blocks

Custom Code Blocks (CCBs) allow you to execute custom Python code within your flows. As your custom code grows in complexity, it becomes increasingly important to test it thoroughly before deploying to production. This guide will walk you through best practices for unit testing your CCBs.

Overview

The Flows SDK provides mock objects that simulate the runtime environment of the Hyperscience Platform, allowing you to test your custom code locally without needing a running instance. This enables rapid development and testing cycles using standard Python testing frameworks like pytest or unittest.

CCBs Backed by Modules

Starting with recent versions of the Flows SDK, CCBs can be backed by entire modules of code rather than just simple functions. This allows you to:

• Organize complex logic into multiple classes and helper functions

• Separate concerns and improve code maintainability

• Write more testable code by isolating business logic

• Reuse code across multiple CCBs

When using the include_module=True parameter, the SDK will serialize your entire module, including all classes and helper functions, making them available at runtime.

Example: Module-Backed CCB

# my_ccb_module.py
class DataProcessor:
    """Helper class for processing data."""
    
    def __init__(self, prefix: str):
        self.prefix = prefix
    
    def process(self, value: str) -> str:
        """Process a value with the configured prefix."""
        return f"{self.prefix}_{value}"

def entrypoint(input_value: str, prefix: str) -> dict:
    """
    Main entrypoint function for the CCB.
    This is what gets called by the Hyperscience Platform.
    """
    processor = DataProcessor(prefix)
    result = processor.process(input_value)
    return {"processed_value": result}

To use this module in a flow:

from flows_sdk.blocks import PythonBlock
from my_ccb_module import entrypoint

# Create a CCB backed by the entire module
my_block = PythonBlock(
    reference_name='my_processor',
    code=entrypoint,
    code_input={
        'input_value': 'hello',
        'prefix': 'output'
    },
    include_module=True  # This includes the entire module
)

Mock Objects for Testing

The Flows SDK provides two main mock objects in flows_sdk.mocks:

MockHsBlockInstance

MockHsBlockInstance simulates the runtime environment and provides implementations for:

• Blob storage: store_blob(), store_blobs(), fetch_blob()

• Logging: log()

• HTTP requests: hs_request(), hs_get(), hs_post()

• Measurements: publish_measurements()

All stored data (blobs, logs, measurements) is kept in memory, making it easy to inspect and assert against in your tests.

MockHsTask

MockHsTask simulates task metadata with properties:

• task_id: Unique identifier for the task

• flow_run_id: Identifier for the flow run

• correlation_id: Correlation identifier across systems

• task_name: Name of the task

• reference_name: Reference name of the block

Testing Strategies

1. Unit Testing Business Logic

Module-based CCBs allow you to separate your business logic from platform-specific code. This allows business logic to be tested independently of the Hyperscience Platform.

# my_business_logic.py
class Calculator:
    """Example for independent logic - no Hyperscience dependencies."""
    
    @staticmethod
    def add(a: int, b: int) -> int:
        return a + b
    
    @staticmethod
    def multiply(a: int, b: int) -> int:
        return a * b

# test_business_logic.py
import unittest
from my_business_logic import Calculator

class TestCalculator(unittest.TestCase):
    
    def test_add(self):
        result = Calculator.add(2, 3)
        self.assertEqual(result, 5)
    
    def test_multiply(self):
        result = Calculator.multiply(4, 5)
        self.assertEqual(result, 20)

2. Testing CCB Entry Points with Mock Objects

For testing the integration with the Hyperscience Platform, use the provided mock objects:

# ccb_with_platform_features.py
from flows_sdk.types import HsBlockInstance, HsTask, StoreBlobRequest

class ReportGenerator:
    """Example of some simple logic that depends on the Hyperscience platform."""
    
    def __init__(self, hs_block_instance: HsBlockInstance):
        self.hs_block_instance = hs_block_instance
    
    def generate_report(self, data: dict) -> str:
        """Generate a report and store it as a blob."""
        report_content = f"Report: {data}".encode('utf-8')
        
        # Log the operation
        self.hs_block_instance.log('Generating report', level=self.hs_block_instance.LogLevel.INFO)
        
        # Store the report as a blob
        blob_response = self.hs_block_instance.store_blob(
            StoreBlobRequest(name='report.txt', content=report_content)
        )
        
        return blob_response.uuid

def entrypoint(data: dict, _hs_block_instance: HsBlockInstance, _hs_task: HsTask) -> dict:
    """CCB entrypoint with system arguments."""
    generator = ReportGenerator(_hs_block_instance)
    blob_uuid = generator.generate_report(data)
    
    return {
        'blob_uuid': blob_uuid,
        'correlation_id': _hs_task.correlation_id
    }

Test this with mock objects:

# test_ccb_with_platform_features.py
import unittest
from flows_sdk.mocks import MockHsBlockInstance, MockHsTask
from ccb_with_platform_features import entrypoint, ReportGenerator

class TestReportGenerator(unittest.TestCase):
    
    def setUp(self):
        """Set up mock objects for each test."""
        self.mock_block = MockHsBlockInstance()
        self.mock_task = MockHsTask(correlation_id='test-correlation-123')
    
    def test_generate_report(self):
        """Test that report generation creates a blob and logs correctly."""
        generator = ReportGenerator(self.mock_block)
        test_data = {'key': 'value'}
        
        blob_uuid = generator.generate_report(test_data)
        
        # Assert blob was created
        self.assertIn(blob_uuid, self.mock_block.blobs)
        
        # Assert blob content is correct
        expected_content = b"Report: {'key': 'value'}"
        self.assertEqual(self.mock_block.blobs[blob_uuid], expected_content)
        
        # Assert logging occurred
        self.assertEqual(len(self.mock_block.log_data), 1)
        self.assertIn('Generating report', self.mock_block.log_data[0])
    
    def test_entrypoint(self):
        """Test the CCB entrypoint function."""
        test_data = {'key': 'value'}
        
        result = entrypoint(
            data=test_data,
            _hs_block_instance=self.mock_block,
            _hs_task=self.mock_task
        )
        
        # Assert return structure
        self.assertIn('blob_uuid', result)
        self.assertIn('correlation_id', result)
        
        # Assert correlation ID is passed through
        self.assertEqual(result['correlation_id'], 'test-correlation-123')
        
        # Assert blob was stored
        blob_uuid = result['blob_uuid']
        self.assertIn(blob_uuid, self.mock_block.blobs)

This allows you to unittest entire CCBs with platform interactions, without having to run them on an actual Hyperscience instance.

3. Testing HTTP Requests

When your CCB makes HTTP requests to the Hyperscience API, you can mock the responses:

# ccb_with_http.py
from flows_sdk.types import HsBlockInstance

def fetch_submission_data(submission_id: str, _hs_block_instance: HsBlockInstance) -> dict:
    """Fetch submission data from the Hyperscience API."""
    response = _hs_block_instance.hs_get(f'/api/v5/submissions/{submission_id}')
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f'Failed to fetch submission: {response.status_code}')

# test_ccb_with_http.py
import unittest
from unittest.mock import Mock
from flows_sdk.mocks import MockHsBlockInstance
from ccb_with_http import fetch_submission_data

class TestFetchSubmissionData(unittest.TestCase):
    
    def setUp(self):
        self.mock_block = MockHsBlockInstance()
    
    def test_successful_fetch(self):
        """Test successful API call."""
        # Create a mock response
        mock_response = Mock()
        mock_response.status_code = 200
        mock_response.json.return_value = {'id': '123', 'status': 'complete'}
        
        # Configure the mock to return this response
        self.mock_block.hs_request.return_value = mock_response
        
        # Call the function
        result = fetch_submission_data('123', self.mock_block)
        
        # Assert the result
        self.assertEqual(result['id'], '123')
        self.assertEqual(result['status'], 'complete')
        
        # Verify the API was called correctly
        self.mock_block.hs_request.assert_called_once_with(
            'GET', '/api/v5/submissions/123'
        )
    
    def test_failed_fetch(self):
        """Test failed API call."""
        # Create a mock response for failure
        mock_response = Mock()
        mock_response.status_code = 404
        
        self.mock_block.hs_request.return_value = mock_response
        
        # Assert that an exception is raised
        with self.assertRaises(Exception) as context:
            fetch_submission_data('999', self.mock_block)
        
        self.assertIn('Failed to fetch submission', str(context.exception))

For more information, see:

• Introduction to Flows SDK

• Examples

• Source Documentation