partial test coverage of pdf extract worker

2 files changed, 70 insertions, 6 deletions

diff --git a/python/sandcrawler/pdf.py b/python/sandcrawler/pdf.py
index b9baed7..cfba679 100644
--- a/python/sandcrawler/pdf.py
+++ b/python/sandcrawler/pdf.py
@@ -1,5 +1,6 @@
 
 import sys
+import datetime
 from io import BytesIO
 from dataclasses import dataclass
 from typing import Optional, Dict, Any
@@ -87,23 +88,24 @@ def process_pdf(blob: bytes, thumb_size=(180,300), thumb_type="JPEG") -> PdfExtr
         pageN = pdf.create_page(n)
         full_text += pageN.text()
     pdf_info = pdf.infos()
-    # Is this actually needed? or does json marshalling work automatically?
-    #for k in pdf_info.keys():
-    #    if isinstance(pdf_info[k], datetime.datetime):
-    #        pdf_info[k] = datetime.datetime.isoformat(pdf_info[k])
+    # TODO: is this actually needed? or does json marshalling work automatically?
+    for k in pdf_info.keys():
+        if isinstance(pdf_info[k], datetime.datetime):
+            pdf_info[k] = datetime.datetime.isoformat(pdf_info[k])
 
     return PdfExtractResult(
         sha1hex=sha1hex,
         file_meta=file_meta,
         status='success',
         error_msg=None,
-        text=full_text,
+        text=full_text or None,
         page0_thumbnail=page0_thumbnail,
-        meta_xml=pdf.metadata,
+        meta_xml=pdf.metadata or None,
         pdf_info=pdf.infos(),
         pdf_extra=dict(
             height=page0rect.height,
             width=page0rect.width,
+            page_count=pdf.pages,
             permanent_id=pdf.pdf_id.permanent_id,
             update_id=pdf.pdf_id.update_id,
             pdf_version=f"{pdf.pdf_version[0]}.{pdf.pdf_version[1]}",
@@ -155,6 +157,7 @@ class PdfExtractBlobWorker(SandcrawlerWorker):
     def process(self, blob, key: Optional[str] = None):
         if not blob:
             return None
+        assert isinstance(blob, bytes)
 
         result = process_pdf(blob)
         if self.thumbnail_sink and result.page0_thumbnail is not None:
diff --git a/python/tests/test_pdf.py b/python/tests/test_pdf.py
new file mode 100644
index 0000000..1ccf85c
--- /dev/null
+++ b/python/tests/test_pdf.py
@@ -0,0 +1,61 @@
+
+import pytest
+import struct
+import responses
+
+from sandcrawler import PdfExtractWorker, PdfExtractBlobWorker, CdxLinePusher, BlackholeSink, WaybackClient
+from sandcrawler.pdf import process_pdf
+from test_wayback import wayback_client, cdx_client
+
+
+FAKE_PDF_BYTES = b"%PDF SOME JUNK" + struct.pack("!q", 112853843)
+
+def test_process_fake_pdf():
+    resp = process_pdf(FAKE_PDF_BYTES)
+    print(resp)
+    assert resp.status == "not-pdf"
+
+def test_process_dummy_pdf():
+    with open('tests/files/dummy.pdf', 'rb') as f:
+        pdf_bytes = f.read()
+    resp = process_pdf(pdf_bytes)
+    assert resp.status == 'success'
+    assert resp.page0_thumbnail is not None
+    assert len(resp.text) > 10
+    assert resp.meta_xml is None
+    assert resp.file_meta['mimetype'] == 'application/pdf'
+    print(resp.pdf_info)
+    print(resp.pdf_extra)
+    assert resp.pdf_info['Author'] == "Evangelos Vlachogiannis"
+    # 595 x 842
+    assert resp.pdf_extra['height'] == 842
+    assert resp.pdf_extra['width'] == 595
+    assert resp.pdf_extra['page_count'] == 1
+
+def test_pdfextract_worker_cdx(wayback_client):
+
+    sink = BlackholeSink()
+    worker = PdfExtractWorker(wayback_client, sink=sink, thumbnail_sink=sink)
+
+    with open('tests/files/example.cdx', 'r') as cdx_file:
+        pusher = CdxLinePusher(
+            worker,
+            cdx_file,
+            filter_http_statuses=[200, 226],
+            filter_mimetypes=['application/pdf'],
+        )
+        pusher_counts = pusher.run()
+        assert pusher_counts['total']
+        assert pusher_counts['pushed'] == 7
+        assert pusher_counts['pushed'] == worker.counts['total']
+
+def test_pdfextract_blob_worker():
+
+    sink = BlackholeSink()
+    worker = PdfExtractBlobWorker(sink=sink, thumbnail_sink=sink)
+
+    with open('tests/files/dummy.pdf', 'rb') as f:
+        pdf_bytes = f.read()
+
+    worker.process(pdf_bytes)
+