Initial commit

tests/test_models.py

@@ -0,0 +1,34 @@
+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")

tests/test_translator.py

@@ -0,0 +1,61 @@
+from pytest import raises, warns
+import warnings
+from bergamot import TranslationModel
+from pybergamot import Translator
+from pybergamot.engine import DirectBergamotModelEngine, ChainBergamotModelsEngine
+
+
+def test_init():
+    Translator()
+
+
+def test__load_model():
+    translator = Translator()
+    with raises(ValueError):
+        translator._load_model("inexistant")
+    with raises(EnvironmentError):
+        translator._load_model("ukr-fin-tiny", download=False)
+    assert isinstance(translator._load_model("en-fr-tiny"), TranslationModel)
+
+
+def test__create_engine():
+    translator = Translator()
+    with raises(ValueError):
+        translator._create_engine("inexistant", "en")
+    with raises(ValueError):
+        translator._create_engine("en", "inexistant")
+    assert isinstance(translator._create_engine("en", "fr"), DirectBergamotModelEngine)
+    assert isinstance(translator._create_engine("fr", "es"), ChainBergamotModelsEngine)
+
+
+def test_load():
+    translator = Translator()
+    with raises(ValueError):
+        translator.load("inexistant")
+    translator.load("es")
+    with warns(RuntimeWarning):
+        translator.load("es")
+    assert len(translator.loaded_languages) == 1 and translator.loaded_languages[0] == "es"
+    assert len(translator._loaded_engines['es']) == 0
+    translator.load("en")
+    assert len(translator.loaded_languages) == 2 and translator.loaded_languages[1] == "en"
+    assert len(translator._loaded_engines['es']) == 1
+    assert len(translator._loaded_engines['en']) == 1
+    translator.load("fr")
+    assert len(translator.loaded_languages) == 3 and translator.loaded_languages[2] == "fr"
+    assert len(translator._loaded_engines['es']) == 2
+    assert len(translator._loaded_engines['en']) == 2
+    assert len(translator._loaded_engines['fr']) == 2
+
+
+def test_translate():
+    translator = Translator()
+    translator.load("en")
+    translator.load("fr")
+    with raises(ValueError):
+        translator.translate("fr", "es", "Salut !")
+    with raises(ValueError):
+        translator.translate("es", "fr", "¡Buenos días!")
+    translator.load("es")
+    assert type(translator.translate("en", "fr", "Hello!")) == str
+    assert type(translator.translate("es", "fr", "¡Buenos días!")) == str