tests/test_mert_encoder.py

import numpy as np
import torch
from unittest.mock import MagicMock

def make_mock_encoder():
    from app.services.mert_encoder import MERTEncoder

    mock_processor = MagicMock()
    mock_processor.return_value = {"input_values": torch.zeros(1, 22050 * 10)}
    mock_processor.sampling_rate = 24000

    fake_hidden = torch.randn(1, 200, 768)
    mock_output = MagicMock()
    mock_output.last_hidden_state = fake_hidden
    mock_model = MagicMock()
    mock_model.return_value = mock_output
    mock_model.eval = MagicMock(return_value=mock_model)

    return MERTEncoder(model=mock_model, processor=mock_processor)

def test_encode_shape():
    enc = make_mock_encoder()
    assert enc.encode(np.random.randn(22050 * 10).astype(np.float32), sr=22050).shape == (768,)

def test_encode_is_finite():
    enc = make_mock_encoder()
    assert np.isfinite(enc.encode(np.random.randn(22050 * 10).astype(np.float32), sr=22050)).all()

def test_encode_batch_shape():
    enc = make_mock_encoder()
    sr = 22050
    chunks = [{"y": np.random.randn(sr * 10).astype(np.float32), "start": i * 5.0, "end": i * 5.0 + 10.0} for i in range(3)]
    assert enc.encode_batch(chunks, sr=sr).shape == (3, 768)