feat: Set up comprehensive Python testing infrastructure with Poetry

[llbbl]

Set up Python Testing Infrastructure

Summary

This PR establishes a comprehensive testing infrastructure for the PureJaxRL project, migrating from a simple requirements.txt to a modern Python development setup using Poetry.

Changes Made

Package Management

• Migrated to Poetry: Created pyproject.toml with full Poetry configuration
• Preserved dependencies: All existing dependencies from requirements.txt maintained
• Python version: Updated to Python 3.9+ to support all dependencies (specifically flashbax)

Testing Infrastructure

• Testing framework: Added pytest as the main testing framework
• Coverage reporting: Integrated pytest-cov with HTML and XML report generation
• Mocking utilities: Added pytest-mock for test doubles
• Test markers: Configured custom markers (unit, integration, slow) for test categorization

Project Structure

tests/
├── __init__.py
├── conftest.py          # Shared pytest fixtures
├── test_setup_validation.py  # Infrastructure validation tests
├── unit/
│   └── __init__.py
└── integration/
    └── __init__.py

Configuration

pytest Configuration (in pyproject.toml)

• Test discovery patterns for finding test files
• Coverage settings with reporting formats
• Strict mode options for better test quality
• Custom markers for test categorization

Coverage Configuration

• Source directories: purejaxrl
• Multiple report formats: terminal, HTML, XML
• Exclusion patterns for test files and virtual environments
• Coverage threshold initially set to 0% (can be increased as tests are added)

Development Tools

• Linting: black, isort, flake8
• Type checking: mypy
• All tools configured in pyproject.toml for consistency

Shared Fixtures (conftest.py)

• temp_dir: Temporary directory for test files
• rng_key: JAX random key for reproducible tests
• mock_config: Mock configuration dictionary
• sample_observation, sample_action, sample_batch: Test data
• setup_jax: JAX configuration for tests
• mock_env_params: Mock environment parameters
• wandb_mock: Mocked Weights & Biases for testing

How to Use

Running Tests

# Install dependencies
poetry install

# Run all tests
poetry run test
# or
poetry run tests

# Run with specific markers
poetry run pytest -m unit
poetry run pytest -m integration
poetry run pytest -m "not slow"

# Run with coverage report
poetry run pytest --cov=purejaxrl --cov-report=html

Writing Tests

1. Place unit tests in tests/unit/
2. Place integration tests in tests/integration/
3. Use fixtures from conftest.py for common test needs
4. Mark tests appropriately:

   @pytest.mark.unit
   def test_something():
       pass
   
   @pytest.mark.integration
   def test_integration():
       pass
   
   @pytest.mark.slow
   def test_slow_operation():
       pass

Validation

All infrastructure has been validated with a comprehensive test suite (test_setup_validation.py) that verifies:

• pytest is running correctly
• Project structure is set up properly
• JAX imports and operations work
• All fixtures are available
• Test markers function correctly
• Coverage tracking is operational
• Dependencies can be imported

Notes

• The coverage threshold is initially set to 0% to allow immediate use. This should be gradually increased as tests are added to the codebase.
• Poetry lock file (poetry.lock) is intentionally not gitignored to ensure reproducible builds.
• Python version requirement is 3.9+ due to dependency constraints (specifically flashbax).
• All pytest options remain available through the poetry run commands.

[Dante-Berth]

Would be possible to use uv instead of poetry ?

[llbbl]

sure if thats what you would preffer.

[llbbl]

Ok switched to uv.

Do you want me update PR description with relevant commands?

[Dante-Berth]

I have no rights about purejaxrl, but i have seen than CleanRL use uv instead of poetry.

[llbbl]

@Dante-Berth oh I thought you were a contributor. Will see what Chris says then I guess. We can always revert.

I'm hopeful that UV will eventually replace pip, but its not in version 1 yet and poetry more mature, so i'm checking for both and defaulting to poetry.