Spaces:
Sleeping
Sleeping
| import unittest | |
| from agent_helpers import ( | |
| build_user_content, | |
| classify_attachment, | |
| cleanup_exact_answer, | |
| is_youtube_question, | |
| is_youtube_visual_question, | |
| ) | |
| class AgentHelperTests(unittest.TestCase): | |
| def test_build_user_content_exposes_task_id_to_tool_agent(self): | |
| content = build_user_content("What is in the attached file?", "abc-123") | |
| self.assertIn("What is in the attached file?", content) | |
| self.assertIn("abc-123", content) | |
| self.assertIn("download_task_file", content) | |
| def test_build_user_content_without_task_id_is_plain_question(self): | |
| self.assertEqual(build_user_content("What is 2 + 2?", None), "What is 2 + 2?") | |
| def test_classify_attachment_prefers_file_suffix(self): | |
| self.assertEqual(classify_attachment("What does it say?", ".mp3"), "audio") | |
| self.assertEqual(classify_attachment("Analyze the table", ".xlsx"), "spreadsheet") | |
| self.assertEqual(classify_attachment("What is shown?", ".png"), "image") | |
| self.assertEqual(classify_attachment("What is the output?", ".py"), "python") | |
| def test_classify_attachment_uses_question_when_suffix_is_missing(self): | |
| self.assertEqual(classify_attachment("Use the attached spreadsheet", ""), "spreadsheet") | |
| self.assertEqual(classify_attachment("What is said in the recording?", ""), "audio") | |
| self.assertEqual(classify_attachment("Review the chess position image", ""), "image") | |
| self.assertEqual(classify_attachment("Run the attached Python code", ""), "python") | |
| self.assertEqual(classify_attachment("Read the attached text file", ""), "text") | |
| def test_youtube_detection(self): | |
| self.assertTrue(is_youtube_question("Watch https://www.youtube.com/watch?v=L1vXCYZAYYM")) | |
| self.assertTrue(is_youtube_question("See https://youtu.be/L1vXCYZAYYM")) | |
| self.assertFalse(is_youtube_question("This mentions video but has no URL")) | |
| def test_youtube_visual_question_detection(self): | |
| self.assertTrue(is_youtube_visual_question("What is the highest number of bird species on camera? https://youtu.be/x")) | |
| self.assertFalse(is_youtube_visual_question("What does Teal'c say in response? https://youtu.be/x")) | |
| def test_cleanup_exact_answer_removes_common_wrappers(self): | |
| self.assertEqual(cleanup_exact_answer("FINAL ANSWER: 519"), "519") | |
| self.assertEqual(cleanup_exact_answer("`b, e`"), "b, e") | |
| self.assertEqual(cleanup_exact_answer("The answer is: Right."), "Right") | |
| if __name__ == "__main__": | |
| unittest.main() | |