import sys import asyncio from types import SimpleNamespace from pageindex import utils def test_remove_fields_keeps_legacy_max_len(): data = { "title": "A long title", "text": "hidden", "nodes": [{"summary": "abcdefghijklmnopqrstuvwxyz"}], } result = utils.remove_fields(data, fields=["text"], max_len=5) assert "text" not in result assert result["title"] == "A lon..." assert result["nodes"][0]["summary"] == "abcde..." def test_create_node_mapping_keeps_legacy_page_ranges(): tree = [ { "node_id": "0001", "title": "Root", "page_index": 1, "nodes": [ {"node_id": "0002", "title": "Child", "page_index": 3, "nodes": []}, ], } ] plain = utils.create_node_mapping(tree) ranged = utils.create_node_mapping(tree, include_page_ranges=True, max_page=8) assert plain["0001"]["title"] == "Root" assert ranged["0001"]["start_index"] == 1 assert ranged["0001"]["end_index"] == 3 assert ranged["0002"]["start_index"] == 3 assert ranged["0002"]["end_index"] == 8 def test_create_node_mapping_prefers_existing_start_end_ranges(): tree = [ { "node_id": "0001", "title": "Root", "start_index": 1, "end_index": 10, "nodes": [ {"node_id": "0002", "title": "Child", "start_index": 3, "end_index": 5}, ], } ] ranged = utils.create_node_mapping(tree, include_page_ranges=True, max_page=12) assert ranged["0001"]["start_index"] == 1 assert ranged["0001"]["end_index"] == 10 assert ranged["0002"]["start_index"] == 3 assert ranged["0002"]["end_index"] == 5 def test_print_tree_keeps_legacy_exclude_fields(capsys): tree = [{"node_id": "0001", "title": "Root", "text": "hidden", "page_index": 1}] utils.print_tree(tree) out = capsys.readouterr().out assert "Root" in out assert "hidden" not in out assert "page_index" not in out def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch): calls = [] class FakeCompletions: async def create(self, **kwargs): calls.append(kwargs) message = SimpleNamespace(content=" answer ") choice = SimpleNamespace(message=message) return SimpleNamespace(choices=[choice]) class FakeAsyncOpenAI: def __init__(self, api_key): self.api_key = api_key self.chat = SimpleNamespace(completions=FakeCompletions()) fake_openai = SimpleNamespace(AsyncOpenAI=FakeAsyncOpenAI) monkeypatch.setitem(sys.modules, "openai", fake_openai) result = asyncio.run(utils.call_llm( "hello", api_key="sk-test", model="gpt-test", temperature=0.2, )) assert result == "answer" assert calls == [{ "model": "gpt-test", "messages": [{"role": "user", "content": "hello"}], "temperature": 0.2, }]