File size: 5,710 Bytes
c72325f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | from __future__ import annotations
import uuid
import pytest
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
import persistence.db as persistence_db
from backend.app.storage import document_pdf_object_key
from persistence.db import Base
from persistence.document_objects import (
DocumentObjectMetadataError,
attach_owned_document_object_metadata,
get_document_object_metadata,
get_owned_document_object_metadata,
)
from persistence.models import Document, User
from persistence.ownership import OwnedResourceNotFoundError
@pytest.fixture
def session_factory(monkeypatch):
engine = create_engine(
"sqlite+pysqlite:///:memory:",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
future=True,
)
factory = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)
Base.metadata.create_all(engine)
monkeypatch.setattr(persistence_db, "SessionLocal", factory)
try:
yield factory
finally:
Base.metadata.drop_all(engine)
engine.dispose()
def _add_document(session_factory, *, user_id: str, document_id: str) -> None:
with session_factory() as session:
session.add(User(id=user_id, email=f"{user_id}@example.test"))
session.add(
Document(
id=document_id,
user_id=user_id,
source_type="upload",
filename="original private report.pdf",
source_hash=f"hash-{document_id}",
cache_key=f"upload:{document_id}",
status="ready",
embedding_model="intfloat/e5-small-v2",
embedding_format="e5-query-passage-v1",
retrieval_mode="faiss_reranker",
reranker_model="cross-encoder/test",
k_initial=20,
k_final=8,
)
)
session.commit()
def test_existing_document_metadata_is_nullable_and_status_is_unchanged(session_factory) -> None:
user_id = str(uuid.uuid4())
document_id = str(uuid.uuid4())
_add_document(session_factory, user_id=user_id, document_id=document_id)
assert get_owned_document_object_metadata(
user_id=user_id,
document_id=document_id,
) is None
with session_factory() as session:
document = session.get(Document, document_id)
assert document is not None
assert document.status == "ready"
assert document.object_key is None
assert document.content_type is None
assert document.byte_size is None
def test_attach_and_load_document_object_metadata(session_factory) -> None:
user_id = str(uuid.uuid4())
document_id = str(uuid.uuid4())
object_key = document_pdf_object_key(document_id)
_add_document(session_factory, user_id=user_id, document_id=document_id)
attached = attach_owned_document_object_metadata(
user_id=user_id,
document_id=document_id,
object_key=object_key,
content_type=" application/pdf ",
byte_size=1234,
)
owned = get_owned_document_object_metadata(user_id=user_id, document_id=document_id)
internal = get_document_object_metadata(document_id=document_id)
assert attached == owned == internal
assert attached.object_key == object_key
assert attached.content_type == "application/pdf"
assert attached.byte_size == 1234
assert attached.user_id == user_id
def test_document_object_metadata_lookup_enforces_ownership(session_factory) -> None:
owner_id = str(uuid.uuid4())
other_user_id = str(uuid.uuid4())
document_id = str(uuid.uuid4())
_add_document(session_factory, user_id=owner_id, document_id=document_id)
with session_factory() as session:
session.add(User(id=other_user_id, email=f"{other_user_id}@example.test"))
session.commit()
with pytest.raises(OwnedResourceNotFoundError, match="Document not found"):
get_owned_document_object_metadata(
user_id=other_user_id,
document_id=document_id,
)
with pytest.raises(OwnedResourceNotFoundError, match="Document not found"):
attach_owned_document_object_metadata(
user_id=other_user_id,
document_id=document_id,
object_key=document_pdf_object_key(document_id),
content_type="application/pdf",
byte_size=1234,
)
@pytest.mark.parametrize("byte_size", [0, -1, True, 1.5])
def test_document_object_metadata_rejects_invalid_byte_size(session_factory, byte_size) -> None:
user_id = str(uuid.uuid4())
document_id = str(uuid.uuid4())
_add_document(session_factory, user_id=user_id, document_id=document_id)
with pytest.raises(DocumentObjectMetadataError, match="positive integer"):
attach_owned_document_object_metadata(
user_id=user_id,
document_id=document_id,
object_key=document_pdf_object_key(document_id),
content_type="application/pdf",
byte_size=byte_size,
)
def test_document_object_metadata_rejects_key_for_another_document(session_factory) -> None:
user_id = str(uuid.uuid4())
document_id = str(uuid.uuid4())
_add_document(session_factory, user_id=user_id, document_id=document_id)
with pytest.raises(DocumentObjectMetadataError, match="does not match"):
attach_owned_document_object_metadata(
user_id=user_id,
document_id=document_id,
object_key=document_pdf_object_key(str(uuid.uuid4())),
content_type="application/pdf",
byte_size=1234,
)
|