from src.investment.exhibit_funcs import Exhibit, get_exhibit_list
from src.prompts.fieldset import FieldSet


def test_add_reimbursement_rows_sets_flag():
    exhibit = Exhibit(exhibit_page="1", exhibit_page_nums=["1"], exhibit_header="H", exhibit_text="text")
    exhibit.add_reimbursement_rows([{"a": 1}], [])
    assert exhibit.has_reimbursements is True
    assert len(exhibit.reimbursement_rows) == 1


def test_set_exhibit_level_and_previous_dynamic_fields():
    prev = Exhibit(exhibit_page="1", exhibit_page_nums=["1"], exhibit_header="H1", exhibit_text="text1")
    dyn_fields = FieldSet()
    prev.dynamic_primary_fields = dyn_fields
    current = Exhibit(exhibit_page="2", exhibit_page_nums=["2"], exhibit_header="H2", exhibit_text="text2", prev_exhibit=prev)

    current.set_exhibit_level_data({"X": "Y"}, FieldSet())
    assert current.exhibit_level_answers["X"] == "Y"
    assert current.get_previous_exhibit_dynamic_fields() is dyn_fields


def test_get_exhibit_list_links_previous():
    text_dict = {"1": "A", "2": "B"}
    mapping = {"1": ["1"], "2": ["2"]}
    headers = {"1": "H1", "2": "H2"}
    exhibits = get_exhibit_list(text_dict, mapping, headers)
    assert len(exhibits) == 2
    assert exhibits[1].prev_exhibit is exhibits[0]