Add ExtractAttachmentsPreprocessor related tests

nbconvert/preprocessors/tests/base.py

@@ -3,6 +3,8 @@
 # Copyright (c) IPython Development Team.
 # Distributed under the terms of the Modified BSD License.
 
+from base64 import b64encode
+
 from nbformat import v4 as nbformat
 
 from ...exporters.exporter import ResourcesDict
@@ -12,7 +14,7 @@
 class PreprocessorTestsBase(TestsBase):
     """Contains test functions preprocessor tests"""
 
-    def build_notebook(self, with_json_outputs=False):
+    def build_notebook(self, with_json_outputs=False, with_attachment=False):
         """Build a notebook in memory for use with preprocessor tests"""
 
         outputs = [
@@ -42,6 +44,19 @@ def build_notebook(self, with_json_outputs=False):
             nbformat.new_markdown_cell(source="$ e $"),
         ]
 
+        if with_attachment:
+            data = b"test"
+            encoded_data = b64encode(data)
+            # this is conversion of bytes to string, not base64 decoding
+            attachments = {"image.png": {"image/png": encoded_data.decode()}}
+            cells.extend(
+                [
+                    nbformat.new_markdown_cell(
+                        source="![image.png](attachment:image.png)", attachments=attachments
+                    )
+                ]
+            )
+
         return nbformat.new_notebook(cells=cells)
 
     def build_resources(self):

nbconvert/preprocessors/tests/test_extractattachments.py

@@ -0,0 +1,87 @@
+"""Tests for the ExtractAttachments preprocessor"""
+
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+import os
+from base64 import b64decode
+
+from ..extractattachments import ExtractAttachmentsPreprocessor
+from .base import PreprocessorTestsBase
+
+
+class TestExtractAttachments(PreprocessorTestsBase):
+    """Contains test functions for extractattachments.py"""
+
+    def build_preprocessor(self):
+        """Make an instance of a preprocessor"""
+        preprocessor = ExtractAttachmentsPreprocessor()
+        preprocessor.enabled = True
+        return preprocessor
+
+    def test_constructor(self):
+        """Can a ExtractAttachmentsPreprocessor be constructed?"""
+        self.build_preprocessor()
+
+    def test_attachment(self):
+        """Test the output of the ExtractAttachmentsPreprocessor"""
+        nb = self.build_notebook(with_attachment=True)
+        res = self.build_resources()
+        preprocessor = self.build_preprocessor()
+        nb, res = preprocessor(nb, res)
+
+        # Check if attachment was extracted.
+        attachments = nb.cells[-1].attachments
+        self.assertIn("image.png", attachments)
+        self.assertIn("image/png", attachments["image.png"])
+        data = attachments["image.png"]["image/png"]
+        # convert to bytes, b64 decode, convert to str
+        data = b64decode(data.encode("utf-8"))
+        self.assertEqual(data, b"test")
+
+        # Verify attachment
+        self.assertIn("image.png", res["outputs"])
+        self.assertEqual(res["outputs"]["image.png"], b"test")
+
+        # Verify cell source changed appropriately
+        src = nb.cells[-1].source
+        self.assertEqual(src, "![image.png](image.png)")
+
+    def test_attachment_with_directory(self):
+        """Test that cell source modifications are correct when files are put in a directory"""
+        nb = self.build_notebook(with_attachment=True)
+        res = self.build_resources()
+        output_dir = "outputs"
+        res["output_files_dir"] = output_dir
+        preprocessor = self.build_preprocessor()
+        nb, res = preprocessor(nb, res)
+
+        # Verify attachment
+        # This can have "\\" separator on Windows
+        file_path = os.path.join("outputs", "image.png")
+        self.assertIn(file_path, res["outputs"])
+
+        # Verify cell source changed appropriately
+        src = nb.cells[-1].source
+        # This shouldn't change on Windows
+        self.assertEqual(src, "![image.png](outputs/image.png)")
+
+    def test_use_separate_dir_config(self):
+        """Test that use_separate_dir and attachment_directory_template work properly"""
+        nb = self.build_notebook(with_attachment=True)
+        res = self.build_resources()
+        res["unique_key"] = "notebook1"  # add notebook name for the folder
+        preprocessor = self.build_preprocessor()
+        preprocessor.use_separate_dir = True
+        preprocessor.attachments_directory_template = "{notebook_name}_custom"
+        nb, res = preprocessor(nb, res)
+
+        # Verify attachment
+        # This can have "\\" separator on Windows
+        file_path = os.path.join("notebook1_custom", "image.png")
+        self.assertIn(file_path, res["attachments"])
+
+        # Verify cell source changed appropriately
+        src = nb.cells[-1].source
+        # This shouldn't change on Windows
+        self.assertEqual(src, "![image.png](notebook1_custom/image.png)")

nbconvert/writers/tests/test_files.py

@@ -19,7 +19,7 @@ def test_basic_output(self):
 
         # Work in a temporary directory.
         with self.create_temp_cwd():
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res: dict = {}
 
             # Create files writer, test output
@@ -54,8 +54,11 @@ def test_extract(self):
 
         # Work in a temporary directory.
         with self.create_temp_cwd():
-            # Create the resoruces dictionary
-            res = {"outputs": {os.path.join("z_files", "a"): b"b"}}
+            # Create the resources dictionary
+            res = {
+                "outputs": {os.path.join("z_files", "a"): b"b"},
+                "attachments": {os.path.join("z_attachments", "c"): b"d"},
+            }
 
             # Create files writer, test output
             writer = FilesWriter()
@@ -73,12 +76,18 @@ def test_extract(self):
                 output = f.read()
                 self.assertEqual(output, "b")
 
+            attachment_file_dest = os.path.join("z_attachments", "c")
+            assert os.path.isfile(attachment_file_dest)
+            with open(attachment_file_dest) as f:
+                content = f.read()
+                self.assertEqual(content, "d")
+
     def test_build_dir(self):
         """Can FilesWriter write to a build dir correctly?"""
 
         # Work in a temporary directory.
         with self.create_temp_cwd():
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res = {"outputs": {os.path.join("z_files", "a"): b"b"}}
 
             # Create files writer, test output
@@ -122,7 +131,7 @@ def test_links(self):
             with open(os.path.join("sub", "c"), "w") as f:
                 f.write("d")
 
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res: dict = {}
 
             # Create files writer, test output
@@ -195,7 +204,7 @@ def test_relpath(self):
             with open(os.path.join("sub", "c"), "w") as f:
                 f.write("d")
 
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res: dict = {}
 
             # Create files writer, test output
@@ -229,7 +238,7 @@ def test_relpath_default(self):
             with open(os.path.join("sub", "c"), "w") as f:
                 f.write("d")
 
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res = {"metadata": {"path": "sub"}}
 
             # Create files writer, test output
@@ -262,7 +271,7 @@ def test_relpath_precedence(self):
             with open(os.path.join("sub", "c"), "w") as f:
                 f.write("d")
 
-            # Create the resoruces dictionary
+            # Create the resources dictionary
             res = {"metadata": {"path": "other_sub"}}
 
             # Create files writer, test output