Add Unit Tests

Add unit tests to AReaL following the project's testing conventions.

When to Use

This skill is triggered when:

• User asks "how do I add tests?"
• User wants to increase test coverage
• User needs to write tests for new functionality
• User wants to understand AReaL testing patterns

Step-by-Step Guide

Step 1: Understand Test Types

AReaL has two main test categories:

areal/tests/test_<module>_<feature>.py

areal/tests/torchrun/run_*.py

Note: All tests are invoked via pytest. Distributed tests use

torchrun

Step 2: Create Test File Structure

Create test file with naming convention:

test_<module>_<feature>.py

import pytest
import torch

# Import the module to test
from areal.dataset.gsm8k import get_gsm8k_sft_dataset
from areal.tests.utils import get_dataset_path  # Optional test utilities
# For mocking tokenizer: from unittest.mock import MagicMock

Step 3: Write Test Functions

Follow Arrange-Act-Assert pattern:

def test_function_under_condition_returns_expected():
    """Test that function returns expected value under condition."""
    # Arrange
    input_data = 5
    expected_output = 10

    # Act
    result = function_under_test(input_data)

    # Assert
    assert result == expected_output

Step 4: Add Pytest Markers and CI Strategy

Use appropriate pytest markers:

@pytest.mark.slow

@pytest.mark.ci

@pytest.mark.slow

@pytest.mark.asyncio

@pytest.mark.skipif(cond, reason=...)

@pytest.mark.parametrize(...)

CI Test Strategy:

• @pytest.mark.slow

  : Excluded from CI by default (CI runs

  pytest -m "not slow"

  )

• @pytest.mark.slow

  +

  @pytest.mark.ci

  : Slow but must run in CI
• No marker: Runs in CI (fast unit tests)

@pytest.mark.asyncio
async def test_async_function():
    result = await async_function()
    assert result == expected

@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_gpu_feature():
    tensor = torch.tensor([1, 2, 3], device="cuda")
    # ... assertions

@pytest.mark.parametrize("batch_size", [1, 4, 16])
def test_with_parameters(batch_size):
    # Parameterized test

@pytest.mark.slow
def test_slow_function():
    # Excluded from CI by default

@pytest.mark.slow
@pytest.mark.ci
def test_slow_but_required_in_ci():
    # Slow but must run in CI

Step 5: Mock Distributed Environment

For unit tests that need distributed mocks:

import torch.distributed as dist

def test_distributed_function(monkeypatch):
    monkeypatch.setattr(dist, "get_rank", lambda: 0)
    monkeypatch.setattr(dist, "get_world_size", lambda: 2)
    result = distributed_function()
    assert result == expected

Step 6: Handle GPU Dependencies

Always skip gracefully when GPU unavailable:

CUDA_AVAILABLE = torch.cuda.is_available()

@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
def test_gpu_function():
    tensor = torch.tensor([1, 2, 3], device="cuda")
    # ... assertions

Key Requirements (Based on testing.md)

Mocking Distributed

• Use

  torch.distributed.fake_pg

  for unit tests
• Mock

  dist.get_rank()

  and

  dist.get_world_size()

  explicitly
• Don't mock internals of FSDP/DTensor

GPU Test Constraints

• Always skip gracefully when GPU unavailable
• Clean up GPU memory:

  torch.cuda.empty_cache()

  in fixtures
• Use smallest possible model/batch for unit tests

Assertions

• Use

  torch.testing.assert_close()

  for tensor comparison
• Specify

  rtol

  /

  atol

  explicitly for numerical tests
• Avoid bare

  assert tensor.equal()

  — no useful error message

Reference Implementations

areal/tests/test_utils.py

areal/tests/test_examples.py

areal/tests/test_fsdp_engine_nccl.py

Common Mistakes

• ❌ Missing test file registration: Ensure file follows

  test_*.py

  naming
• ❌ GPU dependency without skip: Always use

  @pytest.mark.skipif

  for GPU tests
• ❌ Incorrect tensor comparisons: Use

  torch.testing.assert_close()

  not

  assert tensor.equal()

• ❌ Memory leaks in GPU tests: Clean up with

  torch.cuda.empty_cache()

• ❌ Mocking too much: Don't mock FSDP/DTensor internals
• ❌ Unclear test names: Follow

  test_<what>_<condition>_<expected>

  pattern
• ❌ No docstrings: Add descriptive docstrings to test functions

Integration with Other Skills

This skill complements other AReaL development skills:

• After

  /add-dataset

  : Add tests for new dataset loaders
• After

  /add-workflow

  : Add tests for new workflows
• After

  /add-reward

  : Add tests for new reward functions
• With

  planner

  agent: Reference this skill when planning test implementation

Running Tests

# Run specific test file
uv run pytest areal/tests/test_<name>.py

# Skip slow tests (CI default)
uv run pytest -m "not slow"

# Run with verbose output
uv run pytest -v

# Run distributed tests (requires torchrun)
# Note: Usually invoked via pytest test files
torchrun --nproc_per_node=2 areal/tests/torchrun/run_<test>.py