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)]