Spaces:
Running
Running
File size: 5,358 Bytes
f1e8b1a 1bbf278 | 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 | import os
import tempfile
import pytest
from pathlib import Path
from agents.document_parser import (
extract_text,
extract_text_from_bytes,
infer_document_type,
compute_file_hash,
)
def test_plain_text_and_markdown_extraction():
with tempfile.TemporaryDirectory() as tmpdir:
txt_path = os.path.join(tmpdir, "cv.txt")
with open(txt_path, "w", encoding="utf-8") as f:
f.write("# Jane Doe\nSenior Backend Engineer\nPython, Kafka, Distributed Systems.")
md_path = os.path.join(tmpdir, "notes.md")
with open(md_path, "w", encoding="utf-8") as f:
f.write("## Interview Debrief\nStrong grasp of event-driven architectures.")
txt_extracted = extract_text(txt_path)
assert "Jane Doe" in txt_extracted
assert "Distributed Systems" in txt_extracted
md_extracted = extract_text(md_path)
assert "Interview Debrief" in md_extracted
assert "event-driven" in md_extracted
def test_docx_extraction():
from docx import Document
with tempfile.TemporaryDirectory() as tmpdir:
doc_path = os.path.join(tmpdir, "resume.docx")
doc = Document()
doc.add_heading("John Smith", level=1)
doc.add_paragraph("Staff Distributed Systems Engineer with 10 years experience.")
# Add a table
table = doc.add_table(rows=2, cols=2)
table.cell(0, 0).text = "Skill"
table.cell(0, 1).text = "Level"
table.cell(1, 0).text = "Python"
table.cell(1, 1).text = "Expert"
doc.save(doc_path)
extracted = extract_text(doc_path)
assert "John Smith" in extracted
assert "Staff Distributed Systems Engineer" in extracted
assert "Python | Expert" in extracted or "Python" in extracted
def test_pdf_extraction():
import pypdf
with tempfile.TemporaryDirectory() as tmpdir:
pdf_path = os.path.join(tmpdir, "applicant_cv.pdf")
writer = pypdf.PdfWriter()
# Create a blank page with text annotation or draw text
page = writer.add_blank_page(width=300, height=300)
# We can write PDF content stream
with open(pdf_path, "wb") as f:
writer.write(f)
# Even with an empty/blank PDF, extract_text should not crash
extracted = extract_text(pdf_path)
assert isinstance(extracted, str)
def test_legacy_doc_rejection():
with tempfile.TemporaryDirectory() as tmpdir:
doc_path = os.path.join(tmpdir, "old_resume.doc")
with open(doc_path, "wb") as f:
f.write(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1") # OLE header
with pytest.raises(ValueError) as exc:
extract_text(doc_path)
assert "Legacy .doc format is not supported" in str(exc.value)
def test_unsupported_format_rejection():
with tempfile.TemporaryDirectory() as tmpdir:
bad_path = os.path.join(tmpdir, "image.png")
with open(bad_path, "wb") as f:
f.write(b"fake image data")
with pytest.raises(ValueError) as exc:
extract_text(bad_path)
assert "Unsupported file format" in str(exc.value)
def test_file_not_found():
with pytest.raises(FileNotFoundError):
extract_text("non_existent_file_path_12345.pdf")
def test_infer_document_type():
assert infer_document_type("jane_doe_cv.pdf") == "cv"
assert infer_document_type("resume_v2.docx") == "cv"
assert infer_document_type("technical_interview_transcript.txt") == "interview"
assert infer_document_type("coding_assessment_results.md") == "assessment"
assert infer_document_type("system_architecture_rfc.pdf") == "project"
def test_compute_file_hash():
data = b"Hello HireTrace Pipeline"
h1 = compute_file_hash(data)
h2 = compute_file_hash(data)
assert h1 == h2
assert len(h1) == 64 # SHA-256 hex length
def test_extract_text_from_bytes():
txt_bytes = "Software Engineer\n5 years experience".encode("utf-8")
extracted = extract_text_from_bytes(txt_bytes, "profile.txt")
assert "Software Engineer" in extracted
def test_json_extraction():
with tempfile.TemporaryDirectory() as tmpdir:
# Case 1: Dict with text/content key
json_path = os.path.join(tmpdir, "candidate.json")
with open(json_path, "w", encoding="utf-8") as f:
f.write('{"name": "Alice Smith", "text": "Senior Distributed Systems Engineer with 8 years experience."}')
extracted = extract_text(json_path)
assert "Senior Distributed Systems Engineer" in extracted
# Case 2: Dict without explicit text key - pretty formatted JSON
json_path_raw = os.path.join(tmpdir, "skills.json")
with open(json_path_raw, "w", encoding="utf-8") as f:
f.write('{"skills": ["Go", "Kubernetes"], "experience_years": 7}')
extracted_raw = extract_text(json_path_raw)
assert "Kubernetes" in extracted_raw
assert "experience_years" in extracted_raw
# Case 3: extract_text_from_bytes with JSON
json_bytes = b'{"bio": "Staff DevOps Architect", "content": "Specializing in multi-cloud infrastructure."}'
extracted_bytes = extract_text_from_bytes(json_bytes, "bio.json")
assert "Specializing in multi-cloud infrastructure" in extracted_bytes
|