import asyncio import io import json import pytest import pageindex.backend.cloud as cloud_mod from pageindex.backend.cloud import CloudBackend, API_BASE from pageindex.errors import CloudAPIError, DocumentNotFoundError def test_cloud_backend_init(): backend = CloudBackend(api_key="pi-test") assert backend._api_key == "pi-test" assert backend._headers["api_key"] == "pi-test" def test_api_base_url(): assert "pageindex.ai" in API_BASE def test_query_rejects_empty_doc_ids(): backend = CloudBackend(api_key="pi-test") with pytest.raises(ValueError, match="cannot be empty"): backend.query("col", "q", doc_ids=[]) # ── helpers ────────────────────────────────────────────────────────────────── class FakeResponse: def __init__(self, status_code=200, json_data=None, text="", lines=None): self.status_code = status_code self._json = json_data if json_data is not None else {} self.text = text self.content = json.dumps(self._json).encode() if json_data is not None else b"" self._lines = lines or [] def json(self): return self._json def iter_lines(self, decode_unicode=True): yield from self._lines def close(self): pass @pytest.fixture(autouse=True) def _no_sleep(monkeypatch): monkeypatch.setattr(cloud_mod.time, "sleep", lambda *_: None) # ── _request: retry must rewind file objects (empty-upload regression) ────── def test_request_rewinds_file_on_retry(monkeypatch): backend = CloudBackend(api_key="pi-test") payload = b"%PDF-1.4 fake body" fobj = io.BytesIO(payload) uploads = [] def fake_request(method, url, headers=None, timeout=None, **kwargs): # Simulate requests consuming the file body on every attempt. uploads.append(kwargs["files"]["file"].read()) if len(uploads) == 1: return FakeResponse(status_code=500) return FakeResponse(status_code=200, json_data={"doc_id": "d1"}) monkeypatch.setattr(cloud_mod.requests, "request", fake_request) resp = backend._request("POST", "/doc/", files={"file": fobj}, data={}) assert resp == {"doc_id": "d1"} assert uploads[0] == payload # Without seek(0) the retry would upload an empty body. assert uploads[1] == payload # ── list_documents: pagination beyond the API's 100-doc page cap ───────────── def test_list_documents_paginates(monkeypatch): backend = CloudBackend(api_key="pi-test") backend._folder_id_cache["col"] = None # skip folder lookup calls = [] def fake_request(method, path, **kwargs): offset = kwargs["params"]["offset"] calls.append(offset) n = 100 if offset == 0 else 30 return {"documents": [{"id": f"d{offset + i}", "name": ""} for i in range(n)]} monkeypatch.setattr(backend, "_request", fake_request) docs = backend.list_documents("col") assert len(docs) == 130 assert calls == [0, 100] assert docs[-1]["doc_id"] == "d129" # ── folder resolution: plan-limit vs transient errors ──────────────────────── def test_folder_unavailable_warns_and_caches(monkeypatch): backend = CloudBackend(api_key="pi-test") def fake_request(method, path, **kwargs): raise CloudAPIError("Cloud API error 403: upgrade", status_code=403) monkeypatch.setattr(backend, "_request", fake_request) with pytest.warns(UserWarning, match="not available on this plan"): assert backend._get_folder_id("col") is None assert backend._folder_id_cache["col"] is None def test_folder_transient_error_propagates_and_is_not_cached(monkeypatch): backend = CloudBackend(api_key="pi-test") def fake_request(method, path, **kwargs): raise CloudAPIError("Cloud API request failed: connection reset") monkeypatch.setattr(backend, "_request", fake_request) with pytest.raises(CloudAPIError): backend._get_folder_id("col") # A blip must not permanently route documents to the global space. assert "col" not in backend._folder_id_cache # ── doc endpoints: 404 maps to DocumentNotFoundError (local parity) ────────── def test_doc_404_maps_to_document_not_found(monkeypatch): backend = CloudBackend(api_key="pi-test") def fake_request(method, path, **kwargs): raise CloudAPIError("Cloud API error 404: not found", status_code=404) monkeypatch.setattr(backend, "_request", fake_request) with pytest.raises(DocumentNotFoundError): backend.get_document_structure("col", "missing") with pytest.raises(DocumentNotFoundError): backend.delete_document("col", "missing") def test_get_document_include_text_warns(monkeypatch): backend = CloudBackend(api_key="pi-test") monkeypatch.setattr(backend, "_doc_request", lambda *a, **k: {"tree": []}) with pytest.warns(UserWarning, match="include_text is not supported"): backend.get_document("col", "d1", include_text=True) def test_get_page_content_preserves_images(monkeypatch): # Cloud OCR pages carry an `images` list; get_page_content must pass it # through (parity with the local backend and the documented PageContent # shape) so the SDK-prompted UI can render figures — omitting it only when # empty. Real API uses page_index/markdown/images keys. backend = CloudBackend(api_key="pi-test") imgs = [{"path": "fig1.png", "width": 640, "height": 480}] monkeypatch.setattr(backend, "_doc_request", lambda *a, **k: {"result": [ {"page_index": 1, "markdown": "page one", "images": imgs}, {"page_index": 2, "markdown": "page two", "images": []}, # empty -> omitted ]}) out = backend.get_page_content("col", "d1", "1,2") assert out == [ {"page": 1, "content": "page one", "images": imgs}, {"page": 2, "content": "page two"}, ] # ── query_stream: terminal contract and error propagation ─────────────────── def _collect_events(backend, **kwargs): async def _run(): events = [] async for ev in backend.query_stream("col", "q", doc_ids=["d1"], **kwargs): events.append(ev) return events return asyncio.run(_run()) def _sse(block_type, content): return "data: " + json.dumps({ "block_metadata": {"type": block_type}, "choices": [{"delta": {"content": content}}], }) def test_query_stream_emits_answer_done(monkeypatch): backend = CloudBackend(api_key="pi-test") lines = [_sse("text", "Hello "), _sse("text", "world"), "data: [DONE]"] monkeypatch.setattr( cloud_mod.requests, "post", lambda *a, **k: FakeResponse(status_code=200, lines=lines), ) events = _collect_events(backend) assert [e.type for e in events] == ["answer_delta", "answer_delta", "answer_done"] # Same contract as the local backend: answer_done carries the full text. assert events[-1].data == "Hello world" def test_query_stream_http_error_raises(monkeypatch): backend = CloudBackend(api_key="pi-test") monkeypatch.setattr( cloud_mod.requests, "post", lambda *a, **k: FakeResponse(status_code=401, text="unauthorized"), ) async def _run(): async for _ in backend.query_stream("col", "q", doc_ids=["d1"]): pass with pytest.raises(CloudAPIError, match="401"): asyncio.run(_run()) def test_query_stream_connect_failure_raises_instead_of_hanging(monkeypatch): backend = CloudBackend(api_key="pi-test") def fake_post(*a, **k): raise cloud_mod.requests.ConnectionError("dns failure") monkeypatch.setattr(cloud_mod.requests, "post", fake_post) async def _run(): async for _ in backend.query_stream("col", "q", doc_ids=["d1"]): pass with pytest.raises(CloudAPIError, match="request failed"): asyncio.run(_run()) def test_query_uses_long_timeout_and_single_attempt(monkeypatch): """Non-streaming chat completion is non-idempotent and slow: it must get a long timeout and must NOT be retried (each retry re-bills the query).""" backend = CloudBackend(api_key="pi-test") calls = [] def fake_request(method, url, headers=None, **kwargs): calls.append(kwargs) raise cloud_mod.requests.ConnectionError("boom") monkeypatch.setattr(cloud_mod.requests, "request", fake_request) with pytest.raises(CloudAPIError): backend.query("col", "q", doc_ids=["d1"]) assert len(calls) == 1 assert calls[0]["timeout"] == 300 def test_query_stream_early_break_stops_background_thread(monkeypatch): """Consumer breaking early must signal the SSE thread to stop, not let it drain the whole stream in the background.""" import threading import time as _real_time # autouse fixture stubs cloud_mod.time.sleep, not this backend = CloudBackend(api_key="pi-test") drained_all = threading.Event() class SlowResponse: status_code = 200 text = "" def iter_lines(self, decode_unicode=True): for i in range(1000): yield _sse("text", f"chunk{i} ") _real_time.sleep(0.002) # pace so the consumer reliably breaks first drained_all.set() # only reached if the thread was NOT stopped def close(self): pass monkeypatch.setattr(cloud_mod.requests, "post", lambda *a, **k: SlowResponse()) async def _run(): async for _ in backend.query_stream("col", "q", doc_ids=["d1"]): break # consume one event, then abandon asyncio.run(_run()) assert not drained_all.is_set()