pybergamot/tests/test_models.py
2023-11-03 00:36:46 +01:00

35 lines
1 KiB
Python

from pytest import raises
from pybergamot.models import Models
def test_update_models_list():
Models.REPO_FOR_MODEL = {}
Models.AVAILABLE = []
Models.INSTALLED = []
Models.LANGS = []
Models.update_models_list()
assert len(Models.INSTALLED) > 0
assert len(Models.LANGS) > 0
assert len(Models.AVAILABLE) > len(Models.INSTALLED)
assert len(Models.REPO_FOR_MODEL) == len(Models.AVAILABLE)
def test_get_model_languages():
with raises(ValueError):
Models.get_model_languages("inexistant-model")
assert Models.get_model_languages("en-fr-tiny") == ("en", "fr")
assert Models.get_model_languages("ukr-eng-tiny") == ("uk", "en")
def test_get_model_name_for_languages():
assert Models.get_model_name_for_languages("en", "fr") == "en-fr-tiny"
assert Models.get_model_name_for_languages("uk", "en") == "ukr-eng-tiny"
assert Models.get_model_name_for_languages("in", "ex") is None
def test_download():
with raises(ValueError):
Models.download("inexistant-model")
Models.download("en-fr-tiny")