vllm-omni/tests/unit/test_config.py

"""
Unit tests for configuration modules.
"""

import pytest

from vllm_omni.config import (
    DiTCacheConfig,
    DiTCacheTensor,
    DiTConfig,
    OmniStageConfig,
    create_ar_stage_config,
    create_dit_stage_config,
)


class TestOmniStageConfig:
    def test_ar_stage_config_creation(self):
        """Test AR stage configuration creation."""
        config = OmniStageConfig(
            stage_id=0,
            engine_type="AR",
            model_path="test-model",
            input_modalities=["text"],
            output_modalities=["text"],
        )
        assert config.engine_type == "AR"
        assert config.stage_id == 0
        assert config.model_path == "test-model"
        assert config.input_modalities == ["text"]
        assert config.output_modalities == ["text"]

    def test_dit_stage_config_creation(self):
        """Test DiT stage configuration creation."""
        dit_config = DiTConfig(num_inference_steps=50)
        config = OmniStageConfig(
            stage_id=1,
            engine_type="DiT",
            model_path="test-dit-model",
            input_modalities=["text"],
            output_modalities=["image"],
            dit_config=dit_config,
        )
        assert config.engine_type == "DiT"
        assert config.dit_config is not None
        assert config.dit_config.num_inference_steps == 50

    def test_invalid_engine_type(self):
        """Test validation of engine type."""
        with pytest.raises(ValueError):
            OmniStageConfig(
                stage_id=0,
                engine_type="INVALID",
                model_path="test-model",
                input_modalities=["text"],
                output_modalities=["text"],
            )

    def test_empty_modalities(self):
        """Test validation of empty modalities."""
        with pytest.raises(ValueError):
            OmniStageConfig(
                stage_id=0,
                engine_type="AR",
                model_path="test-model",
                input_modalities=[],
                output_modalities=["text"],
            )

        with pytest.raises(ValueError):
            OmniStageConfig(
                stage_id=0,
                engine_type="AR",
                model_path="test-model",
                input_modalities=["text"],
                output_modalities=[],
            )

    def test_dit_requires_config(self):
        """Test that DiT engine requires dit_config."""
        with pytest.raises(ValueError):
            OmniStageConfig(
                stage_id=1,
                engine_type="DiT",
                model_path="test-dit-model",
                input_modalities=["text"],
                output_modalities=["image"],
            )


class TestDiTConfig:
    def test_dit_config_defaults(self):
        """Test DiT configuration with defaults."""
        config = DiTConfig(num_inference_steps=50)
        assert config.use_diffusers is False
        assert config.diffusers_pipeline is None
        assert config.guidance_scale == 7.5
        assert config.height == 512
        assert config.width == 512

    def test_diffusers_config(self):
        """Test DiT configuration with diffusers."""
        config = DiTConfig(
            num_inference_steps=50,
            use_diffusers=True,
            diffusers_pipeline="stable-diffusion",
        )
        assert config.use_diffusers is True
        assert config.diffusers_pipeline == "stable-diffusion"


class TestDiTCacheConfig:
    def test_cache_config_creation(self):
        """Test cache configuration creation."""
        cache_tensors = [
            DiTCacheTensor(
                name="test_tensor",
                shape=[1, 512, 512],
                dtype="float32",
                persistent=True,
            )
        ]

        config = DiTCacheConfig(
            cache_tensors=cache_tensors,
            max_cache_size=1024 * 1024 * 1024,
            cache_strategy="fifo",
        )

        assert len(config.cache_tensors) == 1
        assert config.max_cache_size == 1024 * 1024 * 1024
        assert config.cache_strategy == "fifo"
        assert config.enable_optimization is True


class TestHelperFunctions:
    def test_create_ar_stage_config(self):
        """Test create_ar_stage_config helper function."""
        config = create_ar_stage_config(stage_id=0, model_path="test-model")

        assert config.stage_id == 0
        assert config.engine_type == "AR"
        assert config.model_path == "test-model"
        assert config.input_modalities == ["text"]
        assert config.output_modalities == ["text"]

    def test_create_dit_stage_config(self):
        """Test create_dit_stage_config helper function."""
        dit_config = DiTConfig(num_inference_steps=50)

        config = create_dit_stage_config(
            stage_id=1, model_path="test-dit-model", dit_config=dit_config
        )

        assert config.stage_id == 1
        assert config.engine_type == "DiT"
        assert config.model_path == "test-dit-model"
        assert config.input_modalities == ["text"]
        assert config.output_modalities == ["image"]
        assert config.dit_config is not None