import numpy as np
from models.esm_2_650m import get_embedding


def test_get_embedding_shape_and_type():
    # Example short protein sequence
    sequence = "MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQ"
    embedding = get_embedding(sequence)

    assert isinstance(embedding, np.ndarray)
    assert embedding.ndim == 1
    assert embedding.shape[0] in [1280, 2560]
    assert embedding.dtype == np.float32 or embedding.dtype == np.float64