import unittest
from analogai.deepthink.presentation.postprocessing.postprocessor import postprocess_text
from analogai.deepthink.presentation.enums.pronoun_placeholder import PronounPlaceholder


class TestPostprocessor(unittest.TestCase):
    def setUp(self):
        agent = PronounPlaceholder.AGENT_PLACEHOLDER.value
        respondent = PronounPlaceholder.RESPONDENT_PLACEHOLDER.value
        self.test_cases = [
            {
                "input": f"{respondent} be {agent}'s creator",
                "expected": "You be my creator"
            },
            {
                "input": f"{agent}'s voice sound good",
                "expected": "My voice sound good"
            },
            {
                "input": "if there be cloud it might rain",
                "expected": "If there be cloud it might rain"
            },
            {
                "input": f"what do {agent} think would it rain?",
                "expected": "What do I think would it rain?"
            },
            {
                "input": f"{agent} is cute?",
                "expected": "I am cute?"
            },
            {
                "input": f"{respondent} is {agent}'s creator?",
                "expected": "You are my creator?"
            }
        ]

    def test_postprocess_text(self):
        for test_case in self.test_cases:
            with self.subTest(input_text=test_case["input"]):
                result = postprocess_text(test_case["input"])
                self.assertEqual(result, test_case["expected"])

if __name__ == "__main__":
    unittest.main() 


