import unittest
from TTS.tts.utils.text.punctuation import Punctuation, _DEF_PUNCS

class PunctuationTest(unittest.TestCase):
    def setUp(self):
        self.punctuation = Punctuation()
        self.test_texts = [("This, is my text ... to be striped !! from text?", "This is my text to be striped from text"),
                           ("This, is my text ... to be striped !! from text", "This is my text to be striped from text"),
                           ("This, is my text ... to be striped  from text?", "This is my text to be striped  from text"),
                           ("This, is my text to be striped from text", "This is my text to be striped from text")
                           ]

    def test_get_set_puncs(self):
        self.punctuation.puncs = "-="
        self.assertEqual(self.punctuation.puncs, "-=")

        self.punctuation.puncs = _DEF_PUNCS
        self.assertEqual(self.punctuation.puncs, _DEF_PUNCS)

    def test_strip_punc(self):
        for text, gt in self.test_texts:
            text_striped = self.punctuation.strip(text)
            self.assertEqual(text_striped, gt)

    def test_strip_restore(self):
        for text, gt in self.test_texts:
            text_striped, puncs_map = self.punctuation.strip_to_restore(text)
            text_restored = self.punctuation.restore(text_striped, puncs_map)
            self.assertEqual(' '.join(text_striped), gt)
            self.assertEqual(text_restored[0], text)