"""Tests for array serialization."""
import numpy as np
import pytest
from mcp_common.serialization import deserialize_array, serialize_array
def test_serialize_small_array_inline() -> None:
"""Test small array serialization (inline)."""
arr = np.random.rand(100, 100) # ~80KB
array_id = "test-array-1"
metadata = serialize_array(arr, array_id)
assert metadata["format"] == "inline"
assert metadata["array_id"] == array_id
assert metadata["shape"] == [100, 100]
assert metadata["dtype"] == "float64"
assert "data" in metadata
def test_serialize_large_array_disk() -> None:
"""Test large array serialization (disk)."""
# Create array >100MB
arr = np.random.rand(2000, 2000, 7) # ~224MB
array_id = "test-array-large"
metadata = serialize_array(arr, array_id)
assert metadata["format"] == "disk"
assert metadata["array_id"] == array_id
assert "file_path" in metadata
def test_serialize_deserialize_roundtrip() -> None:
"""Test serialization-deserialization roundtrip."""
arr = np.random.rand(50, 50)
array_id = "test-roundtrip"
# Serialize
metadata = serialize_array(arr, array_id)
# Deserialize
arr_restored = deserialize_array(metadata)
# Verify
np.testing.assert_array_equal(arr, arr_restored)
def test_force_inline() -> None:
"""Test forcing inline serialization."""
# Medium array that would normally use memory format
arr = np.random.rand(1000, 1000, 15) # ~120MB
array_id = "test-force-inline"
metadata = serialize_array(arr, array_id, force_inline=True)
# Should be inline despite size
assert metadata["format"] == "inline"
@pytest.mark.gpu()
def test_serialize_cupy_array() -> None:
"""Test serializing CuPy array."""
try:
import cupy as cp
arr = cp.random.rand(100, 100)
array_id = "test-cupy"
metadata = serialize_array(arr, array_id)
assert metadata["format"] == "inline"
assert metadata["shape"] == [100, 100]
except ImportError:
pytest.skip("CuPy not available")
@pytest.mark.gpu()
def test_deserialize_to_gpu() -> None:
"""Test deserializing to GPU."""
try:
import cupy as cp
arr = np.random.rand(50, 50)
array_id = "test-gpu-deserialize"
metadata = serialize_array(arr, array_id)
arr_gpu = deserialize_array(metadata, use_gpu=True)
assert hasattr(arr_gpu, "__cuda_array_interface__")
np.testing.assert_array_almost_equal(arr, cp.asnumpy(arr_gpu))
except ImportError:
pytest.skip("CuPy not available")
