old-books/_pipeline/tests/test_fetch.py

"""Link files in _incoming/: .url / .txt reading, arXiv rewrite, and the download against a local test server."""
import http.server
import json
import os
import sys
import tempfile
import threading
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, ".."))
sys.path.insert(0, HERE)

import bookio
import fetch
import ingest
import settings

try:
    from helpers import make_pdf
except ImportError:
    make_pdf = None

PDF_BYTES = b""


class Handler(http.server.BaseHTTPRequestHandler):
    hits = []

    def log_message(self, *a):
        pass

    def send(self, code, body=b"", ctype="application/pdf", extra=None):
        self.send_response(code)
        self.send_header("Content-Type", ctype)
        self.send_header("Content-Length", str(len(body)))
        for k, v in (extra or {}).items():
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        Handler.hits.append(self.path)
        if self.path == "/files/Deep-Work.pdf":
            self.send(200, PDF_BYTES)
        elif self.path == "/redirect":
            self.send(302, extra={"Location": "/files/Deep-Work.pdf"})
        elif self.path == "/pdf/2401.01234":                  # no extension, like arxiv.org/pdf/<id>
            self.send(200, PDF_BYTES)
        elif self.path == "/dl?id=7":
            self.send(200, PDF_BYTES, extra={"Content-Disposition": 'attachment; filename="Named By Server.pdf"'})
        elif self.path == "/page":
            self.send(200, b"<!doctype html><html><head><title>Sign in</title></head><body>login</body></html>", "text/html")
        elif self.path == "/cut":
            self.send(200, PDF_BYTES[:200])
        elif self.path == "/empty":
            self.send(200, b"")
        else:
            self.send(404, b"not found", "text/plain")


@unittest.skipIf(make_pdf is None, "reportlab not installed")
class FetchTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        global PDF_BYTES
        d = tempfile.mkdtemp()
        make_pdf(os.path.join(d, "a.pdf"), [["Deep Work", "By Cal Newport"], ["text"]], title="Deep Work")
        PDF_BYTES = open(os.path.join(d, "a.pdf"), "rb").read()
        cls.server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        threading.Thread(target=cls.server.serve_forever, daemon=True).start()
        cls.base = "http://127.0.0.1:%d" % cls.server.server_port

    @classmethod
    def tearDownClass(cls):
        cls.server.shutdown()

    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.d = self.tmp.name
        self.cfg = settings.load()
        self.out = os.path.join(self.d, "out")

    def tearDown(self):
        self.tmp.cleanup()

    def urlfile(self, content, name="Some Link.url"):
        p = os.path.join(self.d, name)
        with open(p, "w", encoding="utf-8") as f:
            f.write(content)
        return p

    def test_download_follows_redirects_and_verifies_pdf(self):
        r = fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/redirect\n"), self.out)
        self.assertEqual(os.path.basename(r["pdf"]), "Deep-Work.pdf")
        self.assertEqual(r["pages"], 2)
        self.assertEqual(open(r["pdf"], "rb").read(), PDF_BYTES)
        self.assertEqual(r["source_url"], self.base + "/redirect")
        self.assertEqual(r["name_hints"], ["Some Link.url"])

    def test_windows_shortcut_format_and_txt(self):
        r = fetch.fetch_from_file(self.cfg, self.urlfile("[InternetShortcut]\r\nURL=%s/files/Deep-Work.pdf\r\nIconIndex=0\r\n" % self.base), self.out)
        self.assertTrue(r["pdf"].endswith("Deep-Work.pdf"))
        r = fetch.fetch_from_file(self.cfg, self.urlfile("# my link\n%s/files/Deep-Work.pdf\n" % self.base, "note.txt"), os.path.join(self.d, "o2"))
        self.assertTrue(os.path.exists(r["pdf"]))

    def test_name_from_url_without_extension_and_from_content_disposition(self):
        r = fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/pdf/2401.01234"), self.out)
        self.assertEqual(os.path.basename(r["pdf"]), "2401.01234.pdf")
        r = fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/dl?id=7"), os.path.join(self.d, "o2"))
        self.assertEqual(os.path.basename(r["pdf"]), "Named By Server.pdf")

    def test_sidecar_json_next_to_the_url_file_is_carried_over(self):
        self.urlfile('{"title": "T", "author": "A"}', "Some Link.json")
        r = fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/files/Deep-Work.pdf"), self.out)
        self.assertEqual(json.load(open(os.path.splitext(r["pdf"])[0] + ".json"))["author"], "A")

    def test_html_page_is_refused_with_a_clear_reason(self):
        with self.assertRaises(fetch.FetchError) as cm:
            fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/page"), self.out)
        self.assertIn("did not return a PDF", str(cm.exception))
        self.assertIn("HTML page", str(cm.exception))
        self.assertEqual(os.listdir(self.out) and [n for n in os.listdir(self.out) if n.endswith(".pdf")], [])

    def test_http_404_cut_off_pdf_and_empty_download(self):
        for path, text in (("/nothing", "HTTP 404 (not found)"), ("/cut", "cannot read it"), ("/empty", "empty")):
            with self.assertRaises(fetch.FetchError) as cm:
                fetch.fetch_from_file(self.cfg, self.urlfile(self.base + path), self.out)
            self.assertIn(text, str(cm.exception), path)

    def test_too_large_is_refused(self):
        self.cfg["fetch"]["max_mb"] = 0.0005          # about 500 bytes; the test PDF is a few KB
        with self.assertRaises(fetch.FetchError) as cm:
            fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/files/Deep-Work.pdf"), self.out)
        self.assertIn("larger than the allowed maximum", str(cm.exception))

    def test_unreachable_host(self):
        with self.assertRaises(fetch.FetchError) as cm:
            fetch.fetch_from_file(self.cfg, self.urlfile("http://127.0.0.1:9/x.pdf"), self.out)
        self.assertIn("refused the connection", str(cm.exception))

    def test_ingest_stores_the_downloaded_pdf_and_records_the_source_url(self):
        r = fetch.fetch_from_file(self.cfg, self.urlfile(self.base + "/files/Deep-Work.pdf"), self.out)
        self.cfg["books_dir"] = os.path.join(self.d, "books")
        self.cfg["work_dir"] = os.path.join(self.d, "work")
        book_dir = ingest.run(self.cfg, r["pdf"], source_url=r["source_url"])
        book = bookio.read_book(book_dir)
        self.assertEqual(book["source_url"], self.base + "/files/Deep-Work.pdf")
        self.assertEqual((book["author"], book["title"]), ("Cal Newport", "Deep Work"))
        self.assertTrue(os.path.exists(os.path.join(book_dir, "Deep-Work.pdf")))


class PureTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()

    def tearDown(self):
        self.tmp.cleanup()

    def read(self, content):
        p = os.path.join(self.tmp.name, "x.url")
        with open(p, "w", encoding="utf-8-sig") as f:
            f.write(content)
        return fetch.read_url_file(p)

    def test_arxiv_abs_is_rewritten(self):
        self.assertEqual(fetch.normalize_url("https://arxiv.org/abs/2401.01234"), "https://arxiv.org/pdf/2401.01234")
        self.assertEqual(fetch.normalize_url("http://www.arxiv.org/abs/math/0301001v2"), "https://arxiv.org/pdf/math/0301001v2")
        self.assertEqual(fetch.normalize_url("https://arxiv.org/pdf/2401.01234"), "https://arxiv.org/pdf/2401.01234")
        self.assertEqual(fetch.normalize_url("https://example.org/a.pdf"), "https://example.org/a.pdf")

    def test_reading_link_files(self):
        self.assertEqual(self.read("https://example.org/a.pdf  \n\n"), "https://example.org/a.pdf")
        self.assertEqual(self.read("[InternetShortcut]\nURL=https://example.org/b.pdf\nIDList=\n"), "https://example.org/b.pdf")
        for bad, why in (("", "empty"), ("ftp://x/y.pdf", "not a web address"), ("see https://a.org/x.pdf", "not a web address"),
                         ("https://a.org/1.pdf\nhttps://a.org/2.pdf", "exactly one")):
            with self.assertRaises(fetch.FetchError) as cm:
                self.read(bad)
            self.assertIn(why, str(cm.exception), bad)

    def test_safe_names(self):
        self.assertEqual(fetch.safe_name("a%20b.pdf"), "a b.pdf")
        self.assertEqual(fetch.safe_name(""), "download.pdf")
        self.assertEqual(fetch.safe_name("2401.01234"), "2401.01234.pdf")
        self.assertEqual(fetch.name_from_response("HTTP/1.1 200 OK\r\n", "https://x.org/files/My%20Book.pdf?x=1"), "My Book.pdf")


if __name__ == "__main__":
    unittest.main()