feat: Set up comprehensive Python testing infrastructure with Poetry

.gitignore

@@ -161,6 +161,13 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
+# VSCode
+.vscode/
+*.code-workspace
+
+# Claude
+.claude/*
+
 outputs/
 logs/
 

pyproject.toml

@@ -0,0 +1,128 @@
+[tool.poetry]
+name = "vista-vwm"
+version = "0.1.0"
+description = "VISTA: Generalizable World Model with Scaling Up Autonomous Driving Data"
+authors = ["Your Name <[email protected]>"]
+readme = "README.md"
+license = "MIT"
+packages = [{include = "vwm"}]
+
+[tool.poetry.dependencies]
+python = "^3.8"
+torch = ">=2.0.1"
+torchvision = "*"
+torchaudio = "*"
+numpy = "*"
+av = "*"
+black = "23.7.0"
+clip = {git = "https://github.com/openai/CLIP.git"}
+deepspeed = "*"
+diffusers = "*"
+ftfy = "*"
+imageio = "*"
+imageio-ffmpeg = "*"
+kornia = "*"
+matplotlib = "*"
+nuscenes-devkit = "*"
+omegaconf = "*"
+opencv-python = "*"
+pandas = "*"
+pillow = "*"
+pytorch-lightning = "*"
+regex = "*"
+scipy = "*"
+tensorboardx = "*"
+timm = "*"
+torchdata = "*"
+torchmetrics = "*"
+tqdm = "*"
+transformers = "*"
+trimesh = "*"
+wandb = "*"
+webdataset = "*"
+xformers = "*"
+einops = "*"
+beartype = "*"
+scikit-image = "*"
+scikit-video = "*"
+h5py = "*"
+lpips = "*"
+safetensors = "*"
+urllib3 = "<1.27,>=1.25.4"
+torchdiffeq = "*"
+controlnet-aux = "*"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.0"
+pytest-cov = "^4.1.0"
+pytest-mock = "^3.11.0"
+
+[tool.poetry.scripts]
+test = "pytest:main"
+tests = "pytest:main"
+
+[tool.pytest.ini_options]
+minversion = "7.0"
+testpaths = ["tests"]
+python_files = ["test_*.py", "*_test.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+    "-ra",
+    "--strict-markers",
+    "--cov=vwm",
+    "--cov-branch",
+    "--cov-report=term-missing:skip-covered",
+    "--cov-report=html:htmlcov",
+    "--cov-report=xml:coverage.xml",
+    "--cov-fail-under=80",
+    "-vv"
+]
+markers = [
+    "unit: Unit tests",
+    "integration: Integration tests",
+    "slow: Tests that take a long time to run"
+]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::PendingDeprecationWarning"
+]
+
+[tool.coverage.run]
+source = ["vwm"]
+branch = true
+parallel = true
+omit = [
+    "*/tests/*",
+    "*/test_*",
+    "*/__init__.py",
+    "*/setup.py",
+    "*/conftest.py"
+]
+
+[tool.coverage.report]
+precision = 2
+show_missing = true
+skip_covered = false
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "if TYPE_CHECKING:",
+    "if __name__ == .__main__.:",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if 0:",
+    "if False:",
+    "@abstractmethod",
+    "@abc.abstractmethod"
+]
+
+[tool.coverage.html]
+directory = "htmlcov"
+
+[tool.coverage.xml]
+output = "coverage.xml"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

tests/conftest.py

@@ -0,0 +1,179 @@
+import os
+import tempfile
+from pathlib import Path
+from typing import Generator
+
+import pytest
+import torch
+from omegaconf import OmegaConf
+
+
+@pytest.fixture
+def temp_dir() -> Generator[Path, None, None]:
+    """Create a temporary directory for test files."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        yield Path(tmpdir)
+
+
+@pytest.fixture
+def mock_config():
+    """Create a mock configuration for testing."""
+    config = OmegaConf.create({
+        "model": {
+            "type": "test_model",
+            "params": {
+                "in_channels": 3,
+                "out_channels": 3,
+                "hidden_dim": 128
+            }
+        },
+        "data": {
+            "batch_size": 4,
+            "num_workers": 0,
+            "dataset": "test_dataset"
+        },
+        "training": {
+            "epochs": 10,
+            "learning_rate": 1e-4,
+            "optimizer": "adam"
+        },
+        "paths": {
+            "data_dir": "/tmp/test_data",
+            "output_dir": "/tmp/test_output"
+        }
+    })
+    return config
+
+
+@pytest.fixture
+def sample_tensor():
+    """Create a sample tensor for testing."""
+    return torch.randn(2, 3, 64, 64)
+
+
+@pytest.fixture
+def sample_video_tensor():
+    """Create a sample video tensor for testing (batch, time, channels, height, width)."""
+    return torch.randn(2, 8, 3, 64, 64)
+
+
+@pytest.fixture
+def mock_dataset_item():
+    """Create a mock dataset item."""
+    return {
+        "image": torch.randn(3, 256, 256),
+        "label": torch.randint(0, 10, (1,)).item(),
+        "metadata": {
+            "filename": "test_image.jpg",
+            "timestamp": 1234567890
+        }
+    }
+
+
+@pytest.fixture
+def mock_model_state():
+    """Create a mock model state dictionary."""
+    return {
+        "encoder.weight": torch.randn(128, 3, 3, 3),
+        "encoder.bias": torch.randn(128),
+        "decoder.weight": torch.randn(3, 128, 3, 3),
+        "decoder.bias": torch.randn(3),
+        "epoch": 5,
+        "global_step": 1000,
+        "optimizer_state": {}
+    }
+
+
+@pytest.fixture
+def device():
+    """Get the appropriate device for testing."""
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+@pytest.fixture
+def mock_checkpoint_path(temp_dir):
+    """Create a mock checkpoint path."""
+    checkpoint_dir = temp_dir / "checkpoints"
+    checkpoint_dir.mkdir(exist_ok=True)
+    return checkpoint_dir / "checkpoint_epoch_5.pt"
+
+
+@pytest.fixture(autouse=True)
+def reset_random_seeds():
+    """Reset random seeds before each test for reproducibility."""
+    import random
+    import numpy as np
+
+    random.seed(42)
+    np.random.seed(42)
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(42)
+        torch.cuda.manual_seed_all(42)
+
+
+@pytest.fixture
+def mock_wandb(monkeypatch):
+    """Mock wandb for testing without actual logging."""
+    class MockWandb:
+        def init(self, *args, **kwargs):
+            return self
+
+        def log(self, *args, **kwargs):
+            pass
+
+        def finish(self):
+            pass
+
+        def watch(self, *args, **kwargs):
+            pass
+
+        class config:
+            @staticmethod
+            def update(*args, **kwargs):
+                pass
+
+    mock = MockWandb()
+    monkeypatch.setattr("wandb.init", mock.init)
+    monkeypatch.setattr("wandb.log", mock.log)
+    monkeypatch.setattr("wandb.finish", mock.finish)
+    monkeypatch.setattr("wandb.watch", mock.watch)
+    monkeypatch.setattr("wandb.config", mock.config)
+    return mock
+
+
+@pytest.fixture
+def mock_env_vars(monkeypatch):
+    """Set up mock environment variables for testing."""
+    test_env_vars = {
+        "CUDA_VISIBLE_DEVICES": "0",
+        "WANDB_MODE": "offline",
+        "PYTHONPATH": "/workspace"
+    }
+    for key, value in test_env_vars.items():
+        monkeypatch.setenv(key, value)
+    return test_env_vars
+
+
+@pytest.fixture
+def capture_stdout(monkeypatch):
+    """Capture stdout for testing print statements."""
+    import io
+    import sys
+
+    captured_output = io.StringIO()
+    monkeypatch.setattr(sys, 'stdout', captured_output)
+
+    yield captured_output
+
+    captured_output.close()
+
+
+def pytest_configure(config):
+    """Configure pytest with custom settings."""
+    config.addinivalue_line(
+        "markers", "gpu: mark test to run only when GPU is available"
+    )
+    config.addinivalue_line(
+        "markers", "requires_data: mark test that requires external data"
+    )