from __future__ import annotations
from pathlib import Path
from app.rag.extract import extract_file, extract_html
def _make_txt(d: Path) -> Path:
p = d / "doc.txt"
p.write_text("El producto pesa 25 kg y tiene garantía de 2 años.", encoding="utf-8")
return p
def _make_docx(d: Path) -> Path:
from docx import Document
doc = Document()
doc.add_paragraph("La garantía es de dos años.")
doc.add_paragraph("Envíos gratis a partir de 199€.")
p = d / "doc.docx"
doc.save(p)
return p
def _make_xlsx(d: Path) -> Path:
from openpyxl import Workbook
wb = Workbook()
ws = wb.active
ws.append(["Modelo", "X100"])
ws.append(["Peso", "25 kg"])
p = d / "doc.xlsx"
wb.save(p)
return p
def _make_pdf(d: Path) -> Path:
from fpdf import FPDF
pdf = FPDF()
pdf.add_page()
pdf.set_font("helvetica", size=12)
pdf.cell(0, 10, text="Manual de usuario TOORX")
p = d / "doc.pdf"
pdf.output(str(p))
return p
def test_extract_txt(tmp_path):
assert "25 kg" in extract_file(_make_txt(tmp_path))
def test_extract_docx(tmp_path):
text = extract_file(_make_docx(tmp_path))
assert "garantía es de dos años" in text
assert "199" in text
def test_extract_xlsx(tmp_path):
text = extract_file(_make_xlsx(tmp_path))
assert "X100" in text
assert "25 kg" in text
def test_extract_pdf(tmp_path):
text = extract_file(_make_pdf(tmp_path))
assert "Manual" in text
assert "TOORX" in text
def test_extract_html_strips_tags_and_scripts():
html = (
"Envíos
Gratis a partir de 199€
"
""
)
text = extract_html(html)
assert "Envíos" in text
assert "Gratis a partir de 199€" in text
assert "evil" not in text
def test_extract_unsupported_raises(tmp_path):
p = tmp_path / "thing.xyz"
p.write_text("hi", encoding="utf-8")
import pytest
with pytest.raises(ValueError):
extract_file(p)
class _Resp:
def __init__(self, *, content=b"", text="", status_code=200, content_type="text/html"):
self.content = content
self.text = text
self.status_code = status_code
self.headers = {"content-type": content_type}
def raise_for_status(self):
if self.status_code >= 400:
raise RuntimeError(f"http {self.status_code}")
def json(self):
import json as _json
return _json.loads(self.text or self.content or "null")
class _FakeClient:
"""Maps URLs to responses; unknown URLs return 404."""
def __init__(self, pages):
self.pages = pages
self.requested = []
async def get(self, url, headers=None):
self.requested.append(url)
return self.pages.get(url, _Resp(status_code=404))
async def test_crawl_follows_internal_links_and_sitemap():
from app.rag.extract import crawl_url
base = "https://shop.example"
pages = {
f"{base}/": _Resp(content=(
f'Inicio
'
f'FAQ Envios'
f'externo'
).encode()),
f"{base}/faq": _Resp(content=b"Garantia de dos anios"),
f"{base}/envios": _Resp(content=b"Envio gratis desde 199"),
# discovered only via sitemap, not linked from any page
f"{base}/oculta": _Resp(content=b"Pagina secreta del sitemap"),
f"{base}/sitemap.xml": _Resp(
text=f"{base}/oculta",
content_type="application/xml",
),
}
client = _FakeClient(pages)
text = await crawl_url(base + "/", client=client)
# every internal page is indexed (landing + linked + sitemap-only)
assert "Garantia de dos anios" in text
assert "Envio gratis desde 199" in text
assert "Pagina secreta del sitemap" in text
# external domain is never fetched
assert "https://otro.com/x" not in client.requested
async def test_crawl_respects_max_pages():
from app.rag.extract import crawl_url
base = "https://big.example"
pages = {f"{base}/sitemap.xml": _Resp(status_code=404)}
# landing links to 10 pages; cap at 3
links = " ".join(f'p{i}' for i in range(10))
pages[f"{base}/"] = _Resp(content=f"{links}".encode())
for i in range(10):
pages[f"{base}/p{i}"] = _Resp(content=f"page {i}".encode())
client = _FakeClient(pages)
await crawl_url(base + "/", max_pages=3, client=client)
# the breadth-first CRAWL is capped at 3 pages. The fixed Shopify enrichment
# probes (/products.json, /policies/*, /pages/*) are additive and don't count
# against max_pages, so exclude them from the crawl-page tally.
extras = ("/products.json", "/policies/", "/pages/")
page_fetches = [u for u in client.requested
if "sitemap" not in u and not any(e in u for e in extras)]
assert len(page_fetches) <= 3
# --- crawl-time SSRF guard (for untrusted start URLs, e.g. the public demo) --
def test_assert_public_host_blocks_private_and_allows_public(monkeypatch):
import pytest
from app.rag import extract
# IP literals are classified directly, no DNS — incl. IPv4-mapped IPv6.
for bad in ("127.0.0.1", "10.0.0.5", "169.254.169.254", "::1", "::ffff:127.0.0.1", ""):
with pytest.raises(extract._BlockedAddress):
extract._assert_public_host(bad)
# A hostname whose resolution includes a private answer is refused...
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34", "10.1.2.3"])
with pytest.raises(extract._BlockedAddress):
extract._assert_public_host("rebind.example")
# ...while an all-public hostname passes (no exception).
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34"])
extract._assert_public_host("good.example")
def _redirect_to_private_transport():
import httpx
def handler(request):
if request.url.host == "shop.test":
return httpx.Response(302, headers={"location": "http://169.254.169.254/secret"})
return httpx.Response(
200, text="CLOUD METADATA TOKEN abc123", headers={"content-type": "text/html"}
)
return httpx.MockTransport(handler)
async def test_crawl_guard_blocks_redirect_into_private_network(monkeypatch):
"""A public start URL that 302-redirects to the cloud-metadata endpoint must
NOT be followed: the internal content never enters the crawl corpus."""
import httpx
from app.rag import extract
monkeypatch.setattr(extract, "_resolve_addrs", lambda h: ["93.184.216.34"])
client = httpx.AsyncClient(transport=_redirect_to_private_transport(), follow_redirects=True)
text, pages = await extract.crawl_site("https://shop.test/", client=client, guard_ssrf=True)
await client.aclose()
assert "CLOUD METADATA TOKEN" not in text
assert pages == [] # the redirect was refused; nothing was fetched
async def test_crawl_without_guard_would_leak_redirect_target(monkeypatch):
"""Control proving the guard above is load-bearing: the SAME redirect WITHOUT
guard_ssrf leaks the internal content (so the block test is not vacuous)."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_redirect_to_private_transport(), follow_redirects=True)
text, _pages = await extract.crawl_site("https://shop.test/", client=client)
await client.aclose()
assert "CLOUD METADATA TOKEN" in text
def _blocked_store_with_reader_transport():
"""Store returns 200 but EMPTY (datacenter IP soft-blocked); the r.jina.ai
reader returns the real rendered text from its own IPs."""
import httpx
def handler(request):
if request.url.host == "r.jina.ai":
return httpx.Response(
200,
text="Title: Tienda Frikadas\n\nVendemos juegos de mesa. Envios en 48h. "
"Devoluciones en 14 dias.",
headers={"content-type": "text/plain"},
)
# the store itself: reachable but serves nothing useful to the crawler
return httpx.Response(200, text="", headers={"content-type": "text/html"})
return httpx.MockTransport(handler)
async def test_remote_fallback_recovers_blocked_store():
"""When a direct crawl comes back empty and remote_fallback=True, crawl_site
recovers the store's text via the reader so the demo still trains."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_blocked_store_with_reader_transport(),
follow_redirects=True)
text, pages = await extract.crawl_site("https://shop.test/", client=client,
remote_fallback=True)
await client.aclose()
assert "juegos de mesa" in text
assert "Devoluciones" in text
assert pages == ["https://shop.test/"] # the home page is recorded
async def test_no_remote_fallback_by_default_leaves_empty():
"""Without remote_fallback the blocked store yields nothing (no surprise
external calls on the normal path)."""
import httpx
from app.rag import extract
client = httpx.AsyncClient(transport=_blocked_store_with_reader_transport(),
follow_redirects=True)
text, _pages = await extract.crawl_site("https://shop.test/", client=client)
await client.aclose()
assert text.strip() == ""
async def test_crawl_enriches_with_shopify_catalog_and_policies():
"""The crawl pulls Shopify's /products.json (catalog) + /policies/* even when
the storefront HTML is thin — so products + SHIPPING/returns are always known."""
import json as _json
from app.rag.extract import crawl_site
base = "https://tienda.example"
pages = {
f"{base}/sitemap.xml": _Resp(status_code=404),
f"{base}/": _Resp(content=b"Bienvenido"),
f"{base}/products.json?limit=250": _Resp(
text=_json.dumps({"products": [{
"title": "Cinta de correr SRX-100", "product_type": "fitness",
"vendor": "Toorx", "body_html": "Plegable, hasta 12 km/h
",
"variants": [{"title": "Default", "price": "499.00"}],
}]}),
content_type="application/json",
),
f"{base}/policies/shipping-policy": _Resp(content=(
"Politica de envios: hacemos envios a toda la peninsula "
"en 24-48h laborables. Gastos de envio gratis desde 50 euros."
).encode()),
}
client = _FakeClient(pages)
text, urls = await crawl_site(base + "/", max_pages=3, client=client)
assert "SRX-100" in text and "499" in text # catalog from products.json
assert "24-48h" in text # shipping policy page
def test_share_download_url_mappings():
from app.rag.extract import _share_download_url
sp = "https://x-my.sharepoint.com/:x:/g/personal/u/IQ123?e=tok"
dl, suf = _share_download_url(sp)
assert dl.endswith("download=1") and "e=tok" in dl and suf == "xlsx" # keeps the share token
word = "https://x-my.sharepoint.com/:w:/g/personal/u/IQ9?e=z"
assert _share_download_url(word)[1] == "docx"
g = "https://drive.google.com/file/d/ABC123/view?usp=sharing"
assert _share_download_url(g)[0] == "https://drive.google.com/uc?export=download&id=ABC123"
sheets = "https://docs.google.com/spreadsheets/d/SHID/edit#gid=0"
assert _share_download_url(sheets)[0].endswith("/export?format=xlsx")
dbx = "https://www.dropbox.com/s/abc/file.pdf?dl=0"
assert _share_download_url(dbx)[0].endswith("dl=1")
assert _share_download_url("https://example.com/page") is None # not a share link
def test_looks_blocked_rejects_walls_keeps_content():
"""The remote-reader cascade must NEVER index a CAPTCHA/'access denied' wall as
if it were the store's real content (that would feed the bot wrong data)."""
from app.rag.extract import _looks_blocked
assert _looks_blocked("Access Denied. " * 30) == ""
assert _looks_blocked("Just a moment... checking your browser") == ""
assert _looks_blocked("Attention Required! Cloudflare") == ""
assert _looks_blocked("hi") == "" # too short
real = "Cintas de correr y bicicletas de spinning Toorx para uso domestico. " * 6
assert _looks_blocked(real) == real
def test_extract_bytes_routes_by_magic_not_extension(monkeypatch):
"""Parse a file by its CONTENT, not its claimed extension — a legacy .xls
(OLE2) downloaded as ".xlsx" (e.g. from SharePoint) must still read."""
from app.rag import extract
monkeypatch.setattr(extract, "_extract_xls", lambda d: "XLS")
monkeypatch.setattr(extract, "_extract_xlsx", lambda d: "XLSX")
monkeypatch.setattr(extract, "_extract_pdf", lambda d: "PDF")
monkeypatch.setattr(extract, "_extract_docx", lambda d: "DOCX")
assert extract.extract_bytes(b"%PDF-1.4\n...", "xlsx") == "PDF" # magic beats suffix
assert extract.extract_bytes(b"PK\x03\x04zip...", "xlsx") == "XLSX"
assert extract.extract_bytes(b"PK\x03\x04zip...", "docx") == "DOCX"
ole2 = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1rest"
assert extract.extract_bytes(ole2, "xlsx") == "XLS" # legacy .xls labeled .xlsx
assert extract.extract_bytes(b"hola mundo", "txt") == "hola mundo"