import torch

import speakerlab
from speakerlab.process.processor import FBank
from speakerlab.utils.builder import dynamic_import


def test_import_and_eres2netv2_forward():
    ERes2NetV2 = dynamic_import("speakerlab.models.eres2net.ERes2NetV2.ERes2NetV2")
    model = ERes2NetV2(feat_dim=80, embedding_size=192, m_channels=64)
    model.eval()
    x = torch.randn(1, 300, 80)
    with torch.no_grad():
        y = model(x)
    assert y.shape == (1, 192)


def test_fbank_forward():
    fbank = FBank(n_mels=80, sample_rate=16000, mean_nor=False)
    wav = torch.randn(1, 16000 * 2)
    feat = fbank(wav)
    assert feat.dim() == 2 and feat.shape[1] == 80
