improve tests

Author: ricanontherun

server/api/views/uploadFile/test_title.py

@@ -1,21 +1,67 @@
 import unittest
-
-import fitz
+from unittest.mock import MagicMock, patch
 
 import title
 
 class TestGenerateTitle(unittest.TestCase):
-  def test_prefers_metadata_title(self):
-    with fitz.open("./testdata/lithium-longterm.pdf") as doc:
-      self.assertEqual("Long-Term Lithium Therapy: Side Effects and Interactions", title.generate_title(doc))
-      pass
-
-  def test_falls_back_to_first_sentence(self):
-    with fitz.open("./testdata/advancespharmaco.pdf") as doc:
-      expected_title = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
-      self.assertEqual(expected_title, title.generate_title(doc))
-
-  def test_australasian_psychiatry(self):
-    with fitz.open("./testdata/creativitystabilizer.pdf") as doc:
-      expected_title = "Impact of mood stabilizers on creativity"
-      self.assertEqual(expected_title, title.generate_title(doc))
+  def test_prefers_metadata_title_if_valid(self):
+    doc = MagicMock()
+    doc.metadata = {"title": "A Study Regarding The Efficacy of Drugs"}
+    self.assertEqual("A Study Regarding The Efficacy of Drugs", title.generate_title(doc))
+
+  def test_falls_back_to_first_page_text_if_metadata_title_is_empty(self):
+    doc = MagicMock()
+    doc.metadata = {"title": ""}
+    doc[0].get_text = MagicMock()
+
+    foo_block = [None] * 7
+    foo_block[4] = "foo"
+    foo_block[6] = 0
+
+    title_block = [None] * 7
+    title_block[4] = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    title_block[6] = 0
+
+    bar_block = [None] * 7
+    bar_block[4] = "bar"
+    bar_block[6] = 0
+    doc[0].get_text.return_value = [foo_block, title_block, bar_block]
+
+    expected_title = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    self.assertEqual(expected_title, title.generate_title(doc))
+
+  def test_falls_back_to_first_page_text_if_metadata_title_does_not_match_regex(self):
+    doc = MagicMock()
+    doc.metadata = {"title": "abcd1234"}
+    doc[0].get_text = MagicMock()
+
+    foo_block = [None] * 7
+    foo_block[4] = "foo"
+    foo_block[6] = 0
+
+    title_block = [None] * 7
+    title_block[4] = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    title_block[6] = 0
+
+    bar_block = [None] * 7
+    bar_block[4] = "bar"
+    bar_block[6] = 0
+    doc[0].get_text.return_value = [foo_block, title_block, bar_block]
+
+    expected_title = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    self.assertEqual(expected_title, title.generate_title(doc))
+
+  @patch("server.api.services.openai_services.openAIServices.openAI")
+  def test_falls_back_to_chatgpt_if_no_title_found(self, mock_openAI):
+    doc = MagicMock()
+    doc.metadata = {"title": None}
+    doc.get_text.return_value = []
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "A Study Regarding The Efficacy of Drugs"
+    mock_openAI.return_value = mock_response
+
+    title.generate_title(doc)
+
+    self.assertTrue(mock_openAI.called)

…uploadFile/testdata/advancespharmaco.pdf …e/testdata/advances-in-mood-disorder.pdfserver/api/views/uploadFile/testdata/advancespharmaco.pdf renamed to server/api/views/uploadFile/testdata/advances-in-mood-disorder.pdf server/api/views/uploadFile/testdata/advancespharmaco.pdf renamed to server/api/views/uploadFile/testdata/advances-in-mood-disorder.pdf

@@ -5,6 +5,7 @@
 from server.api.services.openai_services import openAIServices
 
 # regular expression to match common research white paper titles. Created by Chat-gpt
+# requires at least 3 words, no dates, no version numbers.
 title_regex = re.compile(r'^(?=(?:\b\w+\b[\s:,\-\(\)]*){3,})(?!.*\b(?:19|20)\d{2}\b)(?!.*\bv\d+\b)[A-Za-z0-9][\w\s:,\-\(\)]*[A-Za-z\)]$', re.IGNORECASE)
 
 def generate_title(pdf: fitz.Document) -> str | None:
@@ -13,6 +14,8 @@ def generate_title(pdf: fitz.Document) -> str | None:
         if title_regex.match(document_metadata_title):
             print("suitable title was found in metadata")
             return document_metadata_title.strip()
+        else:
+            print("metadata title did not match regex")
 
     print("Looking for title in first page text")
     first_page = pdf[0]
@@ -23,17 +26,16 @@ def generate_title(pdf: fitz.Document) -> str | None:
         if block[6] == 0 # only include text blocks.
     ]
 
+    # For some reason, extracted PDF text has extra spaces. Collapse them here.
     regex = r"\s{2,}"
     text_blocks = [re.sub(regex, " ", text) for text in text_blocks]
 
-    # replace redundant whitespaces with single space.
     if len(text_blocks) != 0:
         for text in text_blocks:
             if title_regex.match(text):
-                print(f"suitable title was found in first page text {text}")
                 return text
 
-    print("using chatgpt to generate title")
+    print("no suitable title found in first page text. Using GPT-4 to summarize the PDF")
     gpt_title = summarize_pdf(pdf)
     return gpt_title or None