Add the assert_pdf_text() method for asserting text in a pdf

Author: mdmintz

seleniumbase/fixtures/base_case.py

@@ -1856,6 +1856,51 @@ def print_unique_links_with_status_codes(self):
         soup = self.get_beautiful_soup(self.get_page_source())
         page_utils._print_unique_links_with_status_codes(page_url, soup)
 
+    def assert_pdf_text(self, pdf, text, page=None):
+        """ Asserts text in a PDF file.
+            PDF can be either a URL or a file path on the local file system.
+            @Params
+            pdf - The URL or file path of the PDF file.
+            text - The expected text to verify in the PDF.
+            page - The page number of the PDF to use (optional).
+                    If a page number is provided, looks only at that page.
+                        (1 is the first page, 2 is the second page, etc.)
+                    If no page number is provided, looks at all the pages. """
+        import PyPDF2
+        if not pdf.lower().endswith('.pdf'):
+            raise Exception("%s is not a PDF file! (Expecting a .pdf)" % pdf)
+        file_path = None
+        if page_utils.is_valid_url(pdf):
+            if self.get_current_url() != pdf:
+                self.open(pdf)
+            self.download_file(pdf)
+            file_name = pdf.split('/')[-1]
+            file_path = self.get_downloads_folder() + '/' + file_name
+        else:
+            if not os.path.exists(pdf):
+                raise Exception("%s is not a valid URL or file path!" % pdf)
+            file_path = os.path.abspath(pdf)
+        pdf_file_object = open(file_path, "rb")
+        pdf_reader = PyPDF2.PdfFileReader(pdf_file_object, strict=False)
+        num_pages = pdf_reader.numPages
+        if type(page) is int:
+            if page > num_pages:
+                raise Exception("Invalid page number for the PDF!")
+            page = page - 1
+            page_obj = pdf_reader.getPage(page)
+            pdf_page_text = page_obj.extractText()
+            if text not in pdf_page_text:
+                raise Exception("PDF [%s] is missing expected text [%s] on "
+                                "page [%s]!" % (file_path, text, page))
+        else:
+            for page_num in range(num_pages):
+                page_obj = pdf_reader.getPage(page_num)
+                pdf_page_text = page_obj.extractText()
+                if text in pdf_page_text:
+                    return
+            raise Exception("PDF [%s] is missing expected text [%s]!"
+                            "" % (file_path, text))
+
     def create_folder(self, folder):
         """ Creates a folder of the given name if it doesn't already exist. """
         if folder.endswith("/"):