73 lines
3.2 KiB
Python
73 lines
3.2 KiB
Python
from ingestion.segment.io import read_monographs_jsonl, write_monographs_jsonl
|
|
from ingestion.segment.models import Heading, Monograph, SectionSpan
|
|
|
|
|
|
def test_round_trip_preserves_all_fields(tmp_path):
|
|
heading = Heading(text="Chỉ định", physical_page=100, y0=80.0, is_monograph_title=False, section_key="chi_dinh")
|
|
section = SectionSpan(key="chi_dinh", display_name="Chỉ định", heading=heading, text="Điều trị nhiễm HIV.")
|
|
monograph = Monograph(
|
|
drug_id="abacavir", drug_name="ABACAVIR", source_page_range=[100, 101],
|
|
sections={"chi_dinh": section}, atc_codes=["J05AF06"], atc_stated_absent=False,
|
|
)
|
|
path = tmp_path / "monographs.jsonl"
|
|
count = write_monographs_jsonl([monograph], path)
|
|
assert count == 1
|
|
|
|
result = list(read_monographs_jsonl(path))
|
|
assert len(result) == 1
|
|
r = result[0]
|
|
assert r.drug_id == "abacavir"
|
|
assert r.drug_name == "ABACAVIR"
|
|
assert r.source_page_range == [100, 101]
|
|
assert r.atc_codes == ["J05AF06"]
|
|
assert r.sections["chi_dinh"].text == "Điều trị nhiễm HIV."
|
|
assert r.sections["chi_dinh"].heading.section_key == "chi_dinh"
|
|
|
|
|
|
def test_multiple_monographs_round_trip(tmp_path):
|
|
m1 = Monograph(drug_id="a", drug_name="A", source_page_range=[1, 2])
|
|
m2 = Monograph(drug_id="b", drug_name="B", source_page_range=[3, 4])
|
|
path = tmp_path / "monographs.jsonl"
|
|
write_monographs_jsonl([m1, m2], path)
|
|
result = list(read_monographs_jsonl(path))
|
|
assert [r.drug_id for r in result] == ["a", "b"]
|
|
|
|
|
|
def test_empty_write_produces_empty_file(tmp_path):
|
|
path = tmp_path / "monographs.jsonl"
|
|
count = write_monographs_jsonl([], path)
|
|
assert count == 0
|
|
assert list(read_monographs_jsonl(path)) == []
|
|
|
|
|
|
def test_table_blocks_survive_a_write_read_round_trip(tmp_path):
|
|
# the lifted table blocks were being computed in memory and then dropped
|
|
# at the file boundary — 148 blocks existed in the run summary but the
|
|
# JSONL had no "tables" key at all
|
|
from ingestion.segment.models import Heading, Monograph, SectionSpan, TableBlock
|
|
from ingestion.segment.io import read_monographs_jsonl, write_monographs_jsonl
|
|
|
|
heading = Heading(text="Liều lượng và cách dùng", physical_page=339, y0=200.0,
|
|
is_monograph_title=False, section_key="lieu_luong_va_cach_dung")
|
|
m = Monograph(
|
|
drug_id="cefamandol", drug_name="CEFAMANDOL", source_page_range=[338, 340],
|
|
sections={"lieu_luong_va_cach_dung": SectionSpan(
|
|
key="lieu_luong_va_cach_dung", display_name="Liều lượng và cách dùng",
|
|
heading=heading, text="Cách dùng ...")},
|
|
tables=[TableBlock(
|
|
table_id="p339_t0", shape="simple_table", physical_page=339,
|
|
bbox=[40.0, 380.0, 400.0, 620.0],
|
|
section_key="lieu_luong_va_cach_dung",
|
|
text="80 - 50 750 mg - 2 g, 6 giờ/lần.", quarantined=True)],
|
|
)
|
|
path = tmp_path / "m.jsonl"
|
|
write_monographs_jsonl([m], path)
|
|
back = list(read_monographs_jsonl(path))[0]
|
|
assert len(back.tables) == 1
|
|
t = back.tables[0]
|
|
assert t.table_id == "p339_t0"
|
|
assert t.physical_page == 339
|
|
assert t.bbox == [40.0, 380.0, 400.0, 620.0]
|
|
assert t.quarantined is True
|
|
assert "750 mg - 2 g" in t.text
|