Add tests for stress

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+addopts = --ignore=training

requirements-test.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+
+pytest==7.2.0

tests/test_stress.py

@@ -0,0 +1,36 @@
+from ukrainian_tts.stress import sentence_to_stress, stress_with_model
+
+
+def test_stress_table():
+    examples = [
+        #("Бабин біб розцвів у дощ — Буде бабі біб у борщ.\n\nБоронила", "Б+абин б+іб розцв+ів +у д+ощ — Б+уде б+абі б+іб +у б+орщ.\n\nБоронила"),
+        ("Бобер на березі з бобренятами бублики пік.","Боб+ер н+а березі з бобрен+ятами б+ублики п+ік."),
+        (
+            "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району.", 
+            "Кам'ян+ець-Под+ільський - м+істо в Хмельн+ицькій +області Укра+їни, ц+ентр Кам'ян+ець-Под+ільської міськ+ої об'+єднаної територі+альної гром+ади +і Кам'ян+ець-Под+ільського рай+ону."),
+        ("Привіт, як тебе звати?", "Прив+іт, +як теб+е зв+ати?"),
+        ("АННА - український панк-рок гурт", "+АННА - укра+їнський панк-р+ок г+урт"),
+        ("Не тільки в Україні таке може бути.", "Н+е т+ільки в Укра+їні так+е м+оже б+ути."),
+        ("Не тільки в +Укра+їні т+аке може бути.", "Н+е т+ільки в +Укра+їні т+аке м+оже б+ути."),
+        ("два + два", "дв+а + дв+а"),
+        ("Н тльк в крн тк мж бт.", "Н тльк в крн тк мж бт."),
+
+    ]
+    for item in examples:
+        assert sentence_to_stress(item[0]) == item[1]
+
+    examples = [
+        (
+            "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району.", 
+            "к+ам'янець-под+ільський - м+істо в хм+ельницькій обл+асті укра+їни, ц+ентр к+ам'янець-под+ільської м+іської об'+єднаної територі+альної гром+ади +і к+ам'янець-под+ільського рай+ону."
+        ),
+        ("Привіт, як тебе звати?", "прив+іт, +як т+ебе зв+ати?"),
+        ("АННА - український панк-рок гурт", "+анна - укра+їнський п+анк-р+ок г+урт"),
+        ("Не тільки в Україні таке може бути.", "н+е т+ільки в укра+їні т+аке м+оже б+ути."),
+        ("Не тільки в +Укра+їні т+аке може бути.", "н+е т+ільки в ++укра++їні т+ак+е м+оже б+ути."),
+        ("два + два", "дв+а + дв+а"),
+        ("Н тльк в крн тк мж бт.", "н тльк в крн тк мж бт."),
+    ]
+
+    for item in examples:
+        assert stress_with_model(item[0]) == item[1]

ukrainian_tts/stress.py

@@ -74,41 +74,4 @@ def sentence_to_stress(sentence: str, stress_function=stress_dict) -> str:
         for stressed in all_stresses:
             words[stressed] = orig_words[stressed]
         return " ".join(words)
-    return new_stressed
-
-
-if __name__ == "__main__":
-    # TODO: move it to unit tests
-    sentence = "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району."
-    print(sentence_to_stress(sentence))
-    sentence = "Привіт, як тебе звати?"
-    print(sentence_to_stress(sentence))
-    sentence = "АННА - український панк-рок гурт"
-    print(sentence_to_stress(sentence))
-    sentence = "Не тільки в Україні таке може бути."
-    print(sentence_to_stress(sentence))
-    sentence = "Не тільки в +Укра+їні т+аке може бути."
-    print(sentence_to_stress(sentence))
-    sentence = "два + два"
-    print(sentence_to_stress(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(sentence_to_stress(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(sentence_to_stress(sentence))
-
-    sentence = "Кам'янець-Подільський - місто в Хмельницькій області України, центр Кам'янець-Подільської міської об'єднаної територіальної громади і Кам'янець-Подільського району."
-    print(stress_with_model(sentence))
-    sentence = "Привіт, як тебе звати?"
-    print(stress_with_model(sentence))
-    sentence = "АННА - український панк-рок гурт"
-    print(stress_with_model(sentence))
-    sentence = "Не тільки в Україні таке може бути."
-    print(stress_with_model(sentence))
-    sentence = "Не тільки в +Укра+їні т+аке може бути."
-    print(stress_with_model(sentence))
-    sentence = "два + два"
-    print(stress_with_model(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(stress_with_model(sentence))
-    sentence = "Н тльк в крн тк мж бт."
-    print(stress_with_model(sentence))
+    return new_stressed