from tools.graphrag.graphrag import extract_works
from .test_graphrag import make_document, make_chunk


def test_scientific_project_heading_has_research_project_semantics():
    works = extract_works(
        make_document(categories=["vp2024"]),
        [
            make_chunk(
                0,
                ["Vedecký projekt 2024"],
                "Sekcia: Vedecký projekt 2024\nNázov: Overovanie jazykových modelov\n",
            )
        ],
    )
    assert len(works) == 1
    assert works[0]["work_type"] == "research_project"
    assert works[0]["year"] == 2024


def test_bachelor_project_is_separate_from_thesis_category_year():
    works = extract_works(
        make_document(categories=["bp2025"]),
        [
            make_chunk(
                0,
                ["Bakalársky projekt 2024"],
                "Sekcia: Bakalársky projekt 2024\nNázov: Overovanie automatických odpovedí\n",
            )
        ],
    )
    assert [(w["work_type"], w["year"]) for w in works] == [("bachelor_project", 2024)]