Hugging Face's logo

Spaces:
victor34593993
/
flexigo-support-bot
Running

App Files Community
flexigo-support-bot / tests /rag /test_extract.py
victor34593993's picture
victor34593993
free proxy-equivalent cascade (Wayback/CORS/Google/Jina) for bot-blocked sites
3acfe6c verified
Raw
History Blame Contribute Delete
13.7 kB
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 = (
 "<html><body><h1>Envíos</h1><p>Gratis a partir de 199€</p>"
 "<script>evil()</script></body></html>"
 )
 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'<html><body><h1>Inicio</h1>'
 f'<a href="/faq">FAQ</a> <a href="{base}/envios">Envios</a>'
 f'<a href="https://otro.com/x">externo</a></body></html>'
 ).encode()),
 f"{base}/faq": _Resp(content=b"<html><body>Garantia de dos anios</body></html>"),
 f"{base}/envios": _Resp(content=b"<html><body>Envio gratis desde 199</body></html>"),
 # discovered only via sitemap, not linked from any page
 f"{base}/oculta": _Resp(content=b"<html><body>Pagina secreta del sitemap</body></html>"),
 f"{base}/sitemap.xml": _Resp(
 text=f"<urlset><url><loc>{base}/oculta</loc></url></urlset>",
 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'<a href="/p{i}">p{i}</a>' for i in range(10))
 pages[f"{base}/"] = _Resp(content=f"<html><body>{links}</body></html>".encode())
 for i in range(10):
 pages[f"{base}/p{i}"] = _Resp(content=f"<html><body>page {i}</body></html>".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"<html><body>Bienvenido</body></html>"),
 f"{base}/products.json?limit=250": _Resp(
 text=_json.dumps({"products": [{
 "title": "Cinta de correr SRX-100", "product_type": "fitness",
 "vendor": "Toorx", "body_html": "<p>Plegable, hasta 12 km/h</p>",
 "variants": [{"title": "Default", "price": "499.00"}],
 }]}),
 content_type="application/json",
 ),
 f"{base}/policies/shipping-policy": _Resp(content=(
 "<html><body>Politica de envios: hacemos envios a toda la peninsula "
 "en 24-48h laborables. Gastos de envio gratis desde 50 euros.</body></html>"
 ).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"