diff --git a/Tests/test_decompression_bomb.py b/Tests/test_decompression_bomb.py index 53c8bb5e6..1f85b1b9a 100644 --- a/Tests/test_decompression_bomb.py +++ b/Tests/test_decompression_bomb.py @@ -1,37 +1,40 @@ -from tester import * +from helper import unittest, PillowTestCase, tearDownModule from PIL import Image test_file = "Images/lena.ppm" -def test_no_warning_small_file(): - # Implicit assert: no warning. - # A warning would cause a failure. - Image.open(test_file) +class TestDecompressionBomb(PillowTestCase): + 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(): - # Arrange - # Turn limit off - Image.MAX_IMAGE_PIXELS = None - assert_equal(Image.MAX_IMAGE_PIXELS, None) + def test_no_warning_no_limit(self): + # Arrange + # Turn limit off + Image.MAX_IMAGE_PIXELS = None + self.assertEqual(Image.MAX_IMAGE_PIXELS, None) - # Act / Assert - # Implicit assert: no warning. - # A warning would cause a failure. - Image.open(test_file) + # 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) -def test_warning(): - # Arrange - # Set limit to a low, easily testable value - Image.MAX_IMAGE_PIXELS = 10 - assert_equal(Image.MAX_IMAGE_PIXELS, 10) + # Act / Assert + self.assert_warning( + Image.DecompressionBombWarning, + lambda: Image.open(test_file)) - # Act / Assert - assert_warning( - DecompressionBombWarning, - lambda: Image.open(test_file)) +if __name__ == '__main__': + unittest.main() # End of file