tests/test_router.py

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