from pathlib import Path from router import Route, TaskRouter def test_attachment_routes_take_priority(): router = TaskRouter() assert router.route({"question": "Who?", "file_name": "x.mp3"}).route == Route.AUDIO assert router.route({"question": "Who?", "file_name": "x.py"}).route == Route.PYTHON assert ( router.route({"question": "Who?", "file_name": "x.xlsx"}).route == Route.SPREADSHEET ) assert ( router.route({"question": "Who?"}, Path("board.png")).route == Route.IMAGE_CHESS ) def test_content_routes(): router = TaskRouter() assert ( router.route({"question": "Watch https://youtu.be/L1vXCYZAYYM"}).route == Route.YOUTUBE ) assert ( router.route({"question": "|*|a|b|\n|-|-|-|\n|a|a|b|"}).route == Route.MARKDOWN_TABLE ) assert router.route({"question": "Who wrote this article?"}).route == Route.WEB