161 lines
5.4 KiB
Python
161 lines
5.4 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
from io import BytesIO
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from app.auth.deps import get_current_user
|
|
from app.api.routes.documents import router as documents_router
|
|
from app.db.base import Base, get_db
|
|
from app.db.models import Document, DocumentChunk
|
|
from app.db.models import User
|
|
from app.rag.chunker import chunk_blocks
|
|
from app.rag.extractors import ExtractedBlock, extract_document
|
|
from app.rag.ingest import _chunk_payload
|
|
|
|
|
|
def test_extract_text_document_blocks() -> None:
|
|
extracted = extract_document("1. Общие положения\n\nТекст пункта".encode(), "gost.txt")
|
|
|
|
assert extracted.mime_type == "text/plain"
|
|
assert extracted.text.startswith("1. Общие положения")
|
|
assert [block.paragraph_index for block in extracted.blocks] == [0, 2]
|
|
|
|
|
|
def test_extract_docx_document_sections() -> None:
|
|
docx = pytest.importorskip("docx")
|
|
document = docx.Document()
|
|
document.add_heading("1. Термины", level=1)
|
|
document.add_paragraph("Определение из ГОСТ.")
|
|
buffer = BytesIO()
|
|
document.save(buffer)
|
|
|
|
extracted = extract_document(buffer.getvalue(), "terms.docx")
|
|
|
|
assert extracted.mime_type.endswith("wordprocessingml.document")
|
|
assert any(block.section_path == "1. Термины" for block in extracted.blocks)
|
|
assert "Определение из ГОСТ." in extracted.text
|
|
|
|
|
|
def test_extract_pdf_document_pages(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
class FakePage:
|
|
def __init__(self, text: str) -> None:
|
|
self._text = text
|
|
|
|
def extract_text(self) -> str:
|
|
return self._text
|
|
|
|
class FakePdfReader:
|
|
def __init__(self, _: BytesIO) -> None:
|
|
self.pages = [FakePage("Страница 1"), FakePage("Страница 2")]
|
|
|
|
fake_pypdf = types.SimpleNamespace(PdfReader=FakePdfReader)
|
|
monkeypatch.setitem(sys.modules, "pypdf", fake_pypdf)
|
|
|
|
extracted = extract_document(b"%PDF fake", "doc.pdf")
|
|
|
|
assert extracted.mime_type == "application/pdf"
|
|
assert [block.page_start for block in extracted.blocks] == [1, 2]
|
|
assert "Страница 2" in extracted.text
|
|
|
|
|
|
def test_chunk_blocks_keeps_source_metadata() -> None:
|
|
chunks = chunk_blocks(
|
|
[
|
|
ExtractedBlock(text="1. Область применения", page_start=3, page_end=3),
|
|
ExtractedBlock(text="Требование к оформлению.", page_start=3, page_end=3),
|
|
],
|
|
chunk_size=200,
|
|
)
|
|
|
|
assert len(chunks) == 1
|
|
assert chunks[0].page_start == 3
|
|
assert chunks[0].source_ref == "стр. 3"
|
|
assert "Требование" in chunks[0].content
|
|
|
|
|
|
def test_qdrant_payload_contains_citation_fields() -> None:
|
|
doc = Document(id=7, user_id=2, title="ГОСТ", filename="gost.pdf", mime_type="application/pdf")
|
|
row = DocumentChunk(
|
|
id=11,
|
|
document_id=7,
|
|
chunk_index=4,
|
|
content="Пункт 5.3",
|
|
source_filename="gost.pdf",
|
|
mime_type="application/pdf",
|
|
page_start=12,
|
|
page_end=13,
|
|
section_title="5. Требования",
|
|
section_path="5. Требования > 5.3 Оформление",
|
|
paragraph_index=9,
|
|
source_ref="стр. 12-13, 5. Требования > 5.3 Оформление",
|
|
)
|
|
|
|
payload = _chunk_payload(doc, row)
|
|
|
|
assert payload["document_id"] == 7
|
|
assert payload["chunk_id"] == 11
|
|
assert payload["page_start"] == 12
|
|
assert payload["source_ref"].startswith("стр. 12-13")
|
|
|
|
|
|
def test_upload_document_endpoint_stores_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("RAG_ENABLED", "false")
|
|
from app.config import get_settings
|
|
|
|
get_settings.cache_clear()
|
|
engine = create_engine(
|
|
"sqlite://",
|
|
connect_args={"check_same_thread": False},
|
|
poolclass=StaticPool,
|
|
)
|
|
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
db = TestingSessionLocal()
|
|
user = User(username="owner", display_name="Owner", api_token_hash="hash", is_active=True)
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
db.close()
|
|
|
|
app = FastAPI()
|
|
app.include_router(documents_router, prefix="/api/v1")
|
|
|
|
def override_get_db():
|
|
session = TestingSessionLocal()
|
|
try:
|
|
yield session
|
|
finally:
|
|
session.close()
|
|
|
|
def override_current_user():
|
|
session = TestingSessionLocal()
|
|
try:
|
|
return session.get(User, user_id)
|
|
finally:
|
|
session.close()
|
|
|
|
app.dependency_overrides[get_db] = override_get_db
|
|
app.dependency_overrides[get_current_user] = override_current_user
|
|
|
|
with TestClient(app) as client:
|
|
response = client.post(
|
|
"/api/v1/documents/upload",
|
|
files={"file": ("tz.txt", b"1. Scope\n\nRequirement text", "text/plain")},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()["document"]
|
|
assert payload["filename"] == "tz.txt"
|
|
assert payload["chunk_count"] == 1
|
|
assert payload["index_status"] == "stored"
|