mono/artdag/l1/tests/test_frame_compatibility.py

"""
Integration tests for GPU/CPU frame compatibility.

These tests verify that all primitives work correctly with both:
- numpy arrays (CPU frames)
- CuPy arrays (GPU frames)
- GPUFrame wrapper objects

Run with: pytest tests/test_frame_compatibility.py -v
"""

import pytest
import numpy as np
import sys
import os

# Add parent to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Try to import CuPy
try:
    import cupy as cp
    HAS_GPU = True
except ImportError:
    cp = None
    HAS_GPU = False


def create_test_frame(on_gpu=False):
    """Create a test frame (100x100 RGB)."""
    frame = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
    if on_gpu and HAS_GPU:
        return cp.asarray(frame)
    return frame


class MockGPUFrame:
    """Mock GPUFrame for testing without full GPU stack."""
    def __init__(self, data):
        self._data = data

    @property
    def cpu(self):
        if HAS_GPU and hasattr(self._data, 'get'):
            return self._data.get()
        return self._data

    @property
    def gpu(self):
        if HAS_GPU:
            if hasattr(self._data, 'get'):
                return self._data
            return cp.asarray(self._data)
        raise RuntimeError("No GPU available")


class TestColorOps:
    """Test color_ops primitives with different frame types."""

    def test_shift_hsv_numpy(self):
        """shift-hsv should work with numpy arrays."""
        from sexp_effects.primitive_libs.color_ops import prim_shift_hsv
        frame = create_test_frame(on_gpu=False)
        result = prim_shift_hsv(frame, h=30, s=1.2, v=0.9)
        assert isinstance(result, np.ndarray)
        assert result.shape == frame.shape

    @pytest.mark.skipif(not HAS_GPU, reason="No GPU")
    def test_shift_hsv_cupy(self):
        """shift-hsv should work with CuPy arrays."""
        from sexp_effects.primitive_libs.color_ops import prim_shift_hsv
        frame = create_test_frame(on_gpu=True)
        result = prim_shift_hsv(frame, h=30, s=1.2, v=0.9)
        assert isinstance(result, np.ndarray)  # Should return numpy

    def test_shift_hsv_gpuframe(self):
        """shift-hsv should work with GPUFrame objects."""
        from sexp_effects.primitive_libs.color_ops import prim_shift_hsv
        frame = MockGPUFrame(create_test_frame(on_gpu=False))
        result = prim_shift_hsv(frame, h=30, s=1.2, v=0.9)
        assert isinstance(result, np.ndarray)

    def test_invert_numpy(self):
        """invert-img should work with numpy arrays."""
        from sexp_effects.primitive_libs.color_ops import prim_invert_img
        frame = create_test_frame(on_gpu=False)
        result = prim_invert_img(frame)
        assert isinstance(result, np.ndarray)

    def test_adjust_numpy(self):
        """adjust should work with numpy arrays."""
        from sexp_effects.primitive_libs.color_ops import prim_adjust
        frame = create_test_frame(on_gpu=False)
        result = prim_adjust(frame, brightness=10, contrast=1.2)
        assert isinstance(result, np.ndarray)


class TestGeometry:
    """Test geometry primitives with different frame types."""

    def test_rotate_numpy(self):
        """rotate should work with numpy arrays."""
        from sexp_effects.primitive_libs.geometry import prim_rotate
        frame = create_test_frame(on_gpu=False)
        result = prim_rotate(frame, 45)
        assert isinstance(result, np.ndarray)

    def test_scale_numpy(self):
        """scale should work with numpy arrays."""
        from sexp_effects.primitive_libs.geometry import prim_scale
        frame = create_test_frame(on_gpu=False)
        result = prim_scale(frame, 1.5)
        assert isinstance(result, np.ndarray)


class TestBlending:
    """Test blending primitives with different frame types."""

    def test_blend_numpy(self):
        """blend should work with numpy arrays."""
        from sexp_effects.primitive_libs.blending import prim_blend
        frame_a = create_test_frame(on_gpu=False)
        frame_b = create_test_frame(on_gpu=False)
        result = prim_blend(frame_a, frame_b, 0.5)
        assert isinstance(result, np.ndarray)


class TestInterpreterConversion:
    """Test the interpreter's frame conversion."""

    def test_maybe_to_numpy_none(self):
        """_maybe_to_numpy should handle None."""
        from streaming.stream_sexp_generic import StreamInterpreter
        # Create minimal interpreter
        import tempfile
        with tempfile.NamedTemporaryFile(mode='w', suffix='.sexp', delete=False) as f:
            f.write('(stream "test" :fps 30 :width 100 :height 100 (frame frame))')
            f.flush()
            interp = StreamInterpreter(f.name)

        assert interp._maybe_to_numpy(None) is None

    def test_maybe_to_numpy_numpy(self):
        """_maybe_to_numpy should pass through numpy arrays."""
        from streaming.stream_sexp_generic import StreamInterpreter
        import tempfile
        with tempfile.NamedTemporaryFile(mode='w', suffix='.sexp', delete=False) as f:
            f.write('(stream "test" :fps 30 :width 100 :height 100 (frame frame))')
            f.flush()
            interp = StreamInterpreter(f.name)

        frame = create_test_frame(on_gpu=False)
        result = interp._maybe_to_numpy(frame)
        assert result is frame  # Should be same object

    @pytest.mark.skipif(not HAS_GPU, reason="No GPU")
    def test_maybe_to_numpy_cupy(self):
        """_maybe_to_numpy should convert CuPy to numpy."""
        from streaming.stream_sexp_generic import StreamInterpreter
        import tempfile
        with tempfile.NamedTemporaryFile(mode='w', suffix='.sexp', delete=False) as f:
            f.write('(stream "test" :fps 30 :width 100 :height 100 (frame frame))')
            f.flush()
            interp = StreamInterpreter(f.name)

        frame = create_test_frame(on_gpu=True)
        result = interp._maybe_to_numpy(frame)
        assert isinstance(result, np.ndarray)

    def test_maybe_to_numpy_gpuframe(self):
        """_maybe_to_numpy should convert GPUFrame to numpy."""
        from streaming.stream_sexp_generic import StreamInterpreter
        import tempfile
        with tempfile.NamedTemporaryFile(mode='w', suffix='.sexp', delete=False) as f:
            f.write('(stream "test" :fps 30 :width 100 :height 100 (frame frame))')
            f.flush()
            interp = StreamInterpreter(f.name)

        frame = MockGPUFrame(create_test_frame(on_gpu=False))
        result = interp._maybe_to_numpy(frame)
        assert isinstance(result, np.ndarray)


if __name__ == '__main__':
    pytest.main([__file__, '-v'])