Add optional ext keyword argument to set_cover()

.travis.yml

@@ -7,6 +7,7 @@ install:
 - pip install jinja2
 - pip install -e .
 - pip install 'pytest-cov>=2.4.0,<2.6'
+- pip install pytest-mock
 - pip install python-coveralls
 script: py.test --cov=mkepub/ --cov-report=term-missing
 after_success:

mkepub/mkepub.py

@@ -108,9 +108,11 @@ def add_font(self, name, data):
         self.fonts.append(name)
         self._add_file(pathlib.Path('fonts') / name, data)
 
-    def set_cover(self, data):
+    def set_cover(self, data, extension=None):
         """Set the cover image to the given data."""
-        self._cover = 'cover.' + imghdr.what(None, h=data)
+        if not extension:
+            extension = imghdr.what(None, h=data)
+        self._cover = 'cover.' + extension
         self._add_file(pathlib.Path('covers') / self._cover, data)
         self._write('cover.xhtml', 'EPUB/cover.xhtml', cover=self._cover)
 

mkepub/tests/test_mkepub.py

@@ -130,11 +130,47 @@ def test_book_with_dcterms():
     ]
     assert_has_contents(book, 'EPUB/package.opf', expected_strings)
 
-    
+
 def test_mediatype():
     assert mkepub.mkepub.mediatype('file.png') == 'image/png'
     assert mkepub.mkepub.mediatype('file.jpg') == 'image/jpeg'
     assert mkepub.mkepub.mediatype('file.jpeg') == 'image/jpeg'
     assert mkepub.mkepub.mediatype('file.gif') == 'image/gif'
     assert mkepub.mkepub.mediatype('file.svg') == 'image/svg+xml'
 
+
+class TestSetCover:
+    def test_set_cover_auto(self, mocker):
+        mocker.patch("mkepub.Book._add_file")
+        mocker.patch("mkepub.Book._write")
+
+        book = mkepub.Book('SetCover')
+
+        with open('mkepub/tests/cover.jpg', 'rb') as file:
+            data = file.read()
+        book.set_cover(data)
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.jpeg'), data)
+        mkepub.Book._write.assert_called_with(mocker.ANY, mocker.ANY, cover='cover.jpeg')
+
+    def test_set_cover_override(self, mocker):
+        mocker.patch("mkepub.Book._add_file")
+        mocker.patch("mkepub.Book._write")
+
+        book = mkepub.Book('SetCover')
+
+        with open('mkepub/tests/cover.jpg', 'rb') as file:
+            data = file.read()
+        book.set_cover(data, extension='png')
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.png'), data)
+        mkepub.Book._write.assert_called_with(mocker.ANY, mocker.ANY, cover='cover.png')
+
+    def test_set_cover_undetectable(self, mocker):
+        mocker.patch("mkepub.Book._add_file")
+        mocker.patch("mkepub.Book._write")
+
+        book = mkepub.Book('SetCover')
+
+        data = b'<svg></svg>'
+        book.set_cover(data, extension='svg')
+        mkepub.Book._add_file.assert_called_with(pathlib.PosixPath('covers/cover.svg'), data)
+        mkepub.Book._write.assert_called_with(mocker.ANY, mocker.ANY, cover='cover.svg')

setup.py

@@ -22,6 +22,6 @@
         'Programming Language :: Python :: 3.5'],
     packages=['mkepub'],
     package_data={'mkepub': ['templates/*']},
-    tests_require=['epubcheck', 'pytest', 'pytest-cov', 'python-coveralls'],
+    tests_require=['epubcheck', 'pytest', 'pytest-cov', 'pytest-mock', 'python-coveralls'],
     install_requires=['jinja2'],
 )