mirror of
https://github.com/arsvendg/Stirling-PDF.git
synced 2026-07-14 18:44:05 +02:00
James BruntonandGitHub
1264f4cfed
# Description of Changes Change Stirling Engine to support deleting documents automatically. This happens both on user logout and after an amount of time specified by the Java when ingesting a document (allowing for personal documents to have short lifetimes but org documents to be left in the db with no expiry date). Also sets up an [ACL policy](https://en.wikipedia.org/wiki/Access-control_list) for the documents so the database knows which users have access to which documents. This is not fully implemented in the Java, so currently all docs are treated as having a single owner, the uploader, but theoretically when we need to support org storage, we shouldn't need to change the db schema.
371 lines
13 KiB
Python
371 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Iterator
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from stirling.api import app
|
|
from stirling.api.dependencies import get_document_service
|
|
from stirling.documents import Document, DocumentService, SqliteVecStore
|
|
from stirling.models import FileId, PrincipalId, UserId
|
|
|
|
USER = UserId("test-user")
|
|
USER_PRINCIPALS = [PrincipalId("test-user")]
|
|
HEADERS = {"X-User-Id": USER}
|
|
|
|
|
|
class StubEmbedder:
|
|
"""Deterministic embeddings for route tests: no network, no provider needed."""
|
|
|
|
def __init__(self, dim: int = 8) -> None:
|
|
self._dim = dim
|
|
|
|
async def embed_query(self, text: str) -> list[float]:
|
|
h = hash(text) % 1000
|
|
return [(h + i) / 1000.0 for i in range(self._dim)]
|
|
|
|
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
return [await self.embed_query(t) for t in texts]
|
|
|
|
def chunk_and_prepare(
|
|
self,
|
|
text: str,
|
|
source: str = "",
|
|
base_metadata: dict[str, str] | None = None,
|
|
) -> list[Document]:
|
|
from stirling.documents.chunker import chunk_text
|
|
|
|
chunks = chunk_text(text, 100, 10)
|
|
docs = []
|
|
for i, chunk in enumerate(chunks):
|
|
meta = dict(base_metadata) if base_metadata else {}
|
|
meta["source"] = source
|
|
meta["chunk_index"] = str(i)
|
|
doc_id = f"{source}:chunk:{i}" if source else f"chunk:{i}"
|
|
docs.append(Document(id=doc_id, text=chunk, metadata=meta))
|
|
return docs
|
|
|
|
|
|
def _build_service() -> DocumentService:
|
|
return DocumentService(
|
|
embedder=StubEmbedder(), # type: ignore[arg-type]
|
|
store=SqliteVecStore.ephemeral(),
|
|
default_top_k=3,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def service() -> DocumentService:
|
|
return _build_service()
|
|
|
|
|
|
@pytest.fixture
|
|
def client(service: DocumentService) -> Iterator[TestClient]:
|
|
app.dependency_overrides[get_document_service] = lambda: service
|
|
try:
|
|
yield TestClient(app)
|
|
finally:
|
|
app.dependency_overrides.pop(get_document_service, None)
|
|
|
|
|
|
# ── POST /documents ─────────────────────────────────────────────────────
|
|
|
|
|
|
def test_ingest_document_indexes_page_text(client: TestClient, service: DocumentService) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "doc-123",
|
|
"source": "report.pdf",
|
|
"pageText": [
|
|
{"pageNumber": 1, "text": "The introduction covers the main topic."},
|
|
{"pageNumber": 2, "text": "The conclusion summarises the findings."},
|
|
],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["documentId"] == "doc-123"
|
|
assert body["chunksIndexed"] >= 2
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_ingest_document_replaces_existing_content(client: TestClient, service: DocumentService) -> None:
|
|
client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "replace-me",
|
|
"source": "replace-me.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "Original content that existed before."}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
# Second ingest with different content should replace the first entirely
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "replace-me",
|
|
"source": "replace-me.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "New content that replaced the old."}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
results = await service.search("New content", principals=USER_PRINCIPALS, collection=FileId("replace-me"), top_k=5)
|
|
texts = [r.document.text for r in results]
|
|
assert any("New content" in t for t in texts)
|
|
assert not any("Original content" in t for t in texts)
|
|
|
|
|
|
def test_ingest_document_skips_empty_pages(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "mixed",
|
|
"source": "mixed.pdf",
|
|
"pageText": [
|
|
{"pageNumber": 1, "text": " "},
|
|
{"pageNumber": 2, "text": "Real content on page 2."},
|
|
],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["chunksIndexed"] >= 1
|
|
|
|
|
|
def test_ingest_document_with_no_content_returns_zero(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "empty",
|
|
"source": "empty.pdf",
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.json()["chunksIndexed"] == 0
|
|
|
|
|
|
def test_ingest_document_rejects_empty_id(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={"documentId": "", "source": "x.pdf", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_missing_source(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={"documentId": "doc-1", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_empty_source(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={"documentId": "doc-1", "source": "", "pageText": [{"pageNumber": 1, "text": "something"}]},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_non_positive_page_number(client: TestClient) -> None:
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "bad-page",
|
|
"source": "bad-page.pdf",
|
|
"pageText": [{"pageNumber": 0, "text": "something"}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_missing_owner_id(client: TestClient) -> None:
|
|
"""ownerId is required — never derived from the caller. Forgetting it must 422,
|
|
not silently fall back to personal-doc semantics."""
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "no-owner",
|
|
"source": "no-owner.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "something"}],
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_empty_read_principals(client: TestClient) -> None:
|
|
"""readPrincipals is required and must not be empty — every doc needs at least one reader."""
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "no-readers",
|
|
"source": "no-readers.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "something"}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [],
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_ingest_document_rejects_missing_user_header(client: TestClient) -> None:
|
|
"""The route must refuse to write per-user data when the caller didn't identify themselves."""
|
|
response = client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "doc-1",
|
|
"source": "x.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "something"}],
|
|
},
|
|
)
|
|
assert response.status_code == 401
|
|
|
|
|
|
# ── DELETE /documents/{id} ──────────────────────────────────────────────
|
|
|
|
|
|
def test_delete_document_reports_deleted_true_when_existed(client: TestClient) -> None:
|
|
client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "to-delete",
|
|
"source": "to-delete.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
response = client.delete("/api/v1/documents/by-id/to-delete", headers=HEADERS)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"documentId": "to-delete", "deleted": True}
|
|
|
|
|
|
def test_delete_document_is_idempotent(client: TestClient) -> None:
|
|
response = client.delete("/api/v1/documents/by-id/never-existed", headers=HEADERS)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"documentId": "never-existed", "deleted": False}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delete_document_removes_collection(client: TestClient, service: DocumentService) -> None:
|
|
client.post(
|
|
"/api/v1/documents",
|
|
json={
|
|
"documentId": "gone",
|
|
"source": "gone.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "Text."}],
|
|
"ownerId": USER,
|
|
"readPrincipals": [USER],
|
|
"expiresAt": None,
|
|
},
|
|
headers=HEADERS,
|
|
)
|
|
assert await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS)
|
|
client.delete("/api/v1/documents/by-id/gone", headers=HEADERS)
|
|
assert not await service.has_collection(FileId("gone"), principals=USER_PRINCIPALS)
|
|
|
|
|
|
def test_delete_document_rejects_missing_user_header(client: TestClient) -> None:
|
|
response = client.delete("/api/v1/documents/by-id/anything")
|
|
assert response.status_code == 401
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_purge_by_owner_removes_only_callers_collections(client: TestClient, service: DocumentService) -> None:
|
|
"""Logout path: DELETE /api/v1/documents/by-owner purges only the caller's docs."""
|
|
alice = {
|
|
"documentId": "alice-doc",
|
|
"source": "alice.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "alice content"}],
|
|
"ownerId": "alice",
|
|
"readPrincipals": ["alice"],
|
|
"expiresAt": None,
|
|
}
|
|
bob = {
|
|
"documentId": "bob-doc",
|
|
"source": "bob.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "bob content"}],
|
|
"ownerId": "bob",
|
|
"readPrincipals": ["bob"],
|
|
"expiresAt": None,
|
|
}
|
|
client.post("/api/v1/documents", json=alice, headers={"X-User-Id": "alice"})
|
|
client.post("/api/v1/documents", json=bob, headers={"X-User-Id": "bob"})
|
|
|
|
response = client.delete("/api/v1/documents/by-owner", headers={"X-User-Id": "alice"})
|
|
assert response.status_code == 200
|
|
assert response.json() == {"ownerId": "alice", "deleted": 1}
|
|
|
|
# Alice gone, Bob still there.
|
|
assert await service.has_collection(FileId("alice-doc"), principals=[PrincipalId("alice")]) is False
|
|
assert await service.has_collection(FileId("bob-doc"), principals=[PrincipalId("bob")]) is True
|
|
|
|
|
|
def test_purge_by_owner_is_idempotent(client: TestClient) -> None:
|
|
"""Calling purge with no docs is fine — deleted=0 and no error."""
|
|
response = client.delete("/api/v1/documents/by-owner", headers=HEADERS)
|
|
assert response.status_code == 200
|
|
assert response.json() == {"ownerId": USER, "deleted": 0}
|
|
|
|
|
|
def test_purge_by_owner_rejects_missing_user_header(client: TestClient) -> None:
|
|
response = client.delete("/api/v1/documents/by-owner")
|
|
assert response.status_code == 401
|
|
|
|
|
|
def test_delete_document_only_affects_calling_user(client: TestClient) -> None:
|
|
"""Two users with the same document id: one user's delete must not remove the other's."""
|
|
alice_body = {
|
|
"documentId": "shared",
|
|
"source": "shared.pdf",
|
|
"pageText": [{"pageNumber": 1, "text": "x"}],
|
|
"ownerId": "alice",
|
|
"readPrincipals": ["alice"],
|
|
"expiresAt": None,
|
|
}
|
|
bob_body = {**alice_body, "ownerId": "bob", "readPrincipals": ["bob"], "expiresAt": None}
|
|
client.post("/api/v1/documents", json=alice_body, headers={"X-User-Id": "alice"})
|
|
client.post("/api/v1/documents", json=bob_body, headers={"X-User-Id": "bob"})
|
|
|
|
alice_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "alice"})
|
|
assert alice_delete.json() == {"documentId": "shared", "deleted": True}
|
|
|
|
# Bob's copy is still there
|
|
bob_delete = client.delete("/api/v1/documents/by-id/shared", headers={"X-User-Id": "bob"})
|
|
assert bob_delete.json() == {"documentId": "shared", "deleted": True}
|