Skip to content

Conversation

@llbbl
Copy link

@llbbl llbbl commented Jun 28, 2025

Set Up Python Testing Infrastructure

Summary

This PR establishes a comprehensive testing infrastructure for the Samba model project using Poetry for dependency management and pytest as the testing framework. The setup provides a solid foundation for writing and organizing tests with proper tooling and configuration.

Changes Made

1. Package Management

  • Created pyproject.toml with Poetry configuration
  • Migrated dependencies from Dockerfile to Poetry format
  • Separated development dependencies from production dependencies
  • Made CUDA-specific dependencies optional (xformers, flash-attn, causal-conv1d, mamba-ssm) to allow testing in non-CUDA environments

2. Testing Framework Setup

  • pytest as the main testing framework
  • pytest-cov for code coverage reporting
  • pytest-mock for mocking utilities
  • pytest-xdist for parallel test execution
  • Additional dev tools: black, isort, flake8, mypy, pre-commit

3. Testing Configuration

  • Configured pytest settings in pyproject.toml:
    • Test discovery patterns
    • Coverage reporting (HTML, XML, terminal)
    • Custom markers: unit, integration, slow
    • Strict mode and verbose output
  • Coverage configuration:
    • Source directory: lit_gpt
    • Exclusions for test files and generated code
    • HTML and XML report generation

4. Directory Structure

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

5. Shared Fixtures (conftest.py)

  • temp_dir: Temporary directory for test files
  • mock_config: Mock configuration dictionary
  • mock_model_config: Mock Config object
  • device: PyTorch device (CPU/CUDA)
  • sample_tensor: Sample PyTorch tensor
  • mock_tokenizer: Mock tokenizer implementation
  • mock_checkpoint: Mock checkpoint file
  • sample_text_data: Sample text for testing
  • mock_wandb: Mock Weights & Biases integration
  • mock_dataset: Mock dataset file
  • reset_environment: Environment reset between tests
  • capture_stdout: Stdout capture utility

6. Poetry Scripts

  • poetry run test - Run all tests
  • poetry run tests - Alternative command (both work)

7. Additional Setup

  • Updated .gitignore with testing and Claude-related entries
  • Created validation tests to verify the infrastructure

How to Use

Installation

# Install all dependencies
poetry install

# Install only main and dev dependencies (no CUDA)
poetry install --only main,dev

Running Tests

# Run all tests
poetry run test

# Run with coverage
poetry run pytest --cov=lit_gpt

# Run specific test types
poetry run pytest -m unit
poetry run pytest -m integration

# Run slow tests
poetry run pytest --runslow

# Run tests in parallel
poetry run pytest -n auto

Writing Tests

  1. Place unit tests in tests/unit/
  2. Place integration tests in tests/integration/
  3. Use the provided fixtures from conftest.py
  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

Notes

  • CUDA dependencies (xformers, flash-attn, etc.) are commented out in pyproject.toml due to build requirements. These can be installed manually in CUDA-enabled environments.
  • The coverage threshold is currently disabled in the pytest configuration to allow for incremental test development.
  • The validation test for importing lit_gpt is skipped as it requires CUDA dependencies.

Next Steps

  1. Developers can now start writing unit tests for individual components
  2. Integration tests can be added for end-to-end workflows
  3. Consider setting up CI/CD to run tests automatically
  4. Re-enable coverage threshold once sufficient tests are written

- Add pyproject.toml with Poetry package management configuration
- Configure pytest with coverage reporting and custom markers (unit, integration, slow)
- Create test directory structure with proper organization
- Add shared pytest fixtures in conftest.py for common test needs
- Include validation tests to verify infrastructure setup
- Update .gitignore with testing and Claude-related entries
- Configure development scripts for easy test execution (poetry run test/tests)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant