From c5ae491895f60f69a1452d31b0ab04b8b1c4355b Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Mon, 4 Jul 2022 00:48:55 -0400
Subject: [PATCH] fix tests

---
 tts/text/input_parsing/parser.py              |  6 ++-
 .../test_english_phonetization.py             | 51 +++++++++++++++++--
 2 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/tts/text/input_parsing/parser.py b/tts/text/input_parsing/parser.py
index cbdbb87..23854c1 100644
--- a/tts/text/input_parsing/parser.py
+++ b/tts/text/input_parsing/parser.py
@@ -8,7 +8,7 @@ Inspired by and including code from original Tacotron2 by Keith Ito.
 # pyright: strict
 
 import re
-from typing import List, Union
+from typing import List, Union, Any
 
 # Regular expression matching text enclosed in curly braces:
 _curly_re = re.compile(r'(.*?)\{(.+?)\}(.*)')
@@ -16,10 +16,14 @@ _curly_re = re.compile(r'(.*?)\{(.+?)\}(.*)')
 class ArpabetSequence:
     def __init__(self, sequence: str):
         self.sequence = sequence
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, ArpabetSequence) and self.sequence == other.sequence
 
 class GraphemeSequence:
     def __init__(self, sequence: str):
         self.sequence = sequence
+    def __eq__(self, other: Any) -> bool:
+        return isinstance(other, GraphemeSequence) and self.sequence == other.sequence
 
 ParsedSequence = List[Union[ArpabetSequence, GraphemeSequence]]
 
diff --git a/tts/text/phonetization/test_english_phonetization.py b/tts/text/phonetization/test_english_phonetization.py
index deef8d5..5b4b4b5 100644
--- a/tts/text/phonetization/test_english_phonetization.py
+++ b/tts/text/phonetization/test_english_phonetization.py
@@ -1,5 +1,7 @@
 from tts.text.phonetization.cmudict import DEFAULT_CMUDICT_PATH, CMUDict
 from tts.text.phonetization.english_phonetization import EnglishPhonetization
+from tts.text.input_parsing.parser import ArpabetSequence
+from tts.text.input_parsing.parser import GraphemeSequence
 
 class TestEnglishPhonetization:
 
@@ -13,20 +15,59 @@ class TestEnglishPhonetization:
         phoneme_dictionary = CMUDict(DEFAULT_CMUDICT_PATH)
         phonetization = EnglishPhonetization(phoneme_dictionary)
         output = phonetization.try_convert_to_phonemes("this is a test")
-        assert output == ["DH IH1 S", "IH1 Z", "AH0", "T EH1 S T"]
+        assert output == [
+            ArpabetSequence("DH IH1 S"), 
+            ArpabetSequence("IH1 Z"),
+            ArpabetSequence("AH0"), 
+            ArpabetSequence("T EH1 S T"),
+        ]
 
     def test_most_words_in_dictionary(self):
         phoneme_dictionary = CMUDict(DEFAULT_CMUDICT_PATH)
         phonetization = EnglishPhonetization(phoneme_dictionary)
         output = phonetization.try_convert_to_phonemes("nonword is a test")
-        assert output == ["nonword", "IH1 Z", "AH0", "T EH1 S T"]
+        assert output == [
+            GraphemeSequence("nonword"), 
+            ArpabetSequence("IH1 Z"), 
+            ArpabetSequence("AH0"), 
+            ArpabetSequence("T EH1 S T"),
+        ]
 
     def test_punctuation(self):
         phoneme_dictionary = CMUDict(DEFAULT_CMUDICT_PATH)
         phonetization = EnglishPhonetization(phoneme_dictionary)
         output = phonetization.try_convert_to_phonemes("this...")
-        assert output == ["DH IH1 S", "..."]
+        assert output == [
+            ArpabetSequence("DH IH1 S"), 
+            GraphemeSequence("..."),
+        ]
         output = phonetization.try_convert_to_phonemes("this!")
-        assert output == ["DH IH1 S", "!"]
+        assert output == [
+            ArpabetSequence("DH IH1 S"), 
+            GraphemeSequence("!"),
+        ]
         output = phonetization.try_convert_to_phonemes("this?")
-        assert output == ["DH IH1 S", "?"]
+        assert output == [
+            ArpabetSequence("DH IH1 S"), 
+            GraphemeSequence("?"),
+        ]
+
+    def test_extra_spaces(self):
+        phoneme_dictionary = CMUDict(DEFAULT_CMUDICT_PATH)
+        phonetization = EnglishPhonetization(phoneme_dictionary)
+        output = phonetization.try_convert_to_phonemes("    this    is    a     test   ")
+        assert output == [
+            ArpabetSequence("DH IH1 S"), 
+            ArpabetSequence("IH1 Z"),
+            ArpabetSequence("AH0"), 
+            ArpabetSequence("T EH1 S T"),
+        ]
+
+    def test_multi_non_word_extra_spaces(self):
+        phoneme_dictionary = CMUDict(DEFAULT_CMUDICT_PATH)
+        phonetization = EnglishPhonetization(phoneme_dictionary)
+        output = phonetization.try_convert_to_phonemes("    nonword    anothernonword   ")
+        assert output == [
+            GraphemeSequence("nonword"), 
+            GraphemeSequence("anothernonword"), 
+        ]
\ No newline at end of file