diff --git a/pdf_reports/tools.py b/pdf_reports/tools.py
index 56b3ccf..cbf5431 100644
--- a/pdf_reports/tools.py
+++ b/pdf_reports/tools.py
@@ -12,6 +12,8 @@
 import datetime
 import textwrap
 
+from matplotlib.axes import Axes
+
 
 def dataframe_to_html(
     dataframe,
@@ -157,8 +159,7 @@ def figure_data(fig, size=None, fmt="png", bbox_inches="tight", **kwargs):
     **kwargs
       Any other option of Matplotlib's figure.savefig() method.
     """
-    if "AxesSubplot" in str(fig.__class__):
-        # A matplotlib axis was provided: take its containing figure.
+    if isinstance(fig, Axes):
         fig = fig.figure
     output = BytesIO()
     original_size = fig.get_size_inches()
diff --git a/tests/test_basics.py b/tests/test_basics.py
index bc1cc54..12257ae 100644
--- a/tests/test_basics.py
+++ b/tests/test_basics.py
@@ -5,7 +5,7 @@
 
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
-from PyPDF2 import PdfFileReader
+from PyPDF2 import PdfReader
 
 
 def test_pug_to_html():
@@ -25,8 +25,8 @@ def test_basics(tmpdir):
     html = pug_to_html(template_path, title="Summary of your order")
     write_report(html, pdf_path)
     with open(pdf_path, "rb") as f:
-        reader = PdfFileReader(f)
-        assert reader.numPages == 2
+        reader = PdfReader(f)
+        assert len(reader.pages) == 2
 
     pdf_data = write_report(html)
     assert len(pdf_data) > 10000
@@ -49,8 +49,8 @@ def test_with_plots_and_tables(tmpdir):
     html = pug_to_html(template_path, dataframe=dataframe)
     write_report(html, pdf_path)
     with open(pdf_path, "rb") as f:
-        reader = PdfFileReader(f)
-        assert reader.numPages == 2
+        reader = PdfReader(f)
+        assert len(reader.pages) == 2
 
     pdf_data = write_report(html)
     assert len(pdf_data) > 10000