Pillow/Tests/test_decompression_bomb.py
2017-06-21 02:36:11 -07:00

59 lines
1.6 KiB
Python

from helper import unittest, PillowTestCase, hopper
from PIL import Image
TEST_FILE = "Tests/images/hopper.ppm"
ORIGINAL_LIMIT = Image.MAX_IMAGE_PIXELS
class TestDecompressionBomb(PillowTestCase):
def tearDown(self):
Image.MAX_IMAGE_PIXELS = ORIGINAL_LIMIT
def test_no_warning_small_file(self):
# Implicit assert: no warning.
# A warning would cause a failure.
Image.open(TEST_FILE)
def test_no_warning_no_limit(self):
# Arrange
# Turn limit off
Image.MAX_IMAGE_PIXELS = None
self.assertIsNone(Image.MAX_IMAGE_PIXELS)
# Act / Assert
# Implicit assert: no warning.
# A warning would cause a failure.
Image.open(TEST_FILE)
def test_warning(self):
# Arrange
# Set limit to a low, easily testable value
Image.MAX_IMAGE_PIXELS = 10
self.assertEqual(Image.MAX_IMAGE_PIXELS, 10)
# Act / Assert
self.assert_warning(Image.DecompressionBombWarning,
lambda: Image.open(TEST_FILE))
class TestDecompressionCrop(PillowTestCase):
def setUp(self):
self.src = hopper()
Image.MAX_IMAGE_PIXELS = self.src.height * self.src.width
def tearDown(self):
Image.MAX_IMAGE_PIXELS = ORIGINAL_LIMIT
def testEnlargeCrop(self):
# Crops can extend the extents, therefore we should have the
# same decompression bomb warnings on them.
box = (0, 0, self.src.width * 2, self.src.height * 2)
self.assert_warning(Image.DecompressionBombWarning,
lambda: self.src.crop(box))
if __name__ == '__main__':
unittest.main()