Skip to content

Commit

Permalink
Use monkeypatch
Browse files Browse the repository at this point in the history
  • Loading branch information
radarhere committed Jan 21, 2025
1 parent cf7dd2f commit bf94539
Show file tree
Hide file tree
Showing 10 changed files with 94 additions and 128 deletions.
21 changes: 10 additions & 11 deletions Tests/check_png_dos.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,26 +3,25 @@
import zlib
from io import BytesIO

import pytest

from PIL import Image, ImageFile, PngImagePlugin

TEST_FILE = "Tests/images/png_decompression_dos.png"


def test_ignore_dos_text() -> None:
ImageFile.LOAD_TRUNCATED_IMAGES = True
def test_ignore_dos_text(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)

try:
im = Image.open(TEST_FILE)
with Image.open(TEST_FILE) as im:
im.load()
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False

assert isinstance(im, PngImagePlugin.PngImageFile)
for s in im.text.values():
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"
assert isinstance(im, PngImagePlugin.PngImageFile)
for s in im.text.values():
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"

for s in im.info.values():
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"
for s in im.info.values():
assert len(s) < 1024 * 1024, "Text chunk larger than 1M"


def test_dos_text() -> None:
Expand Down
19 changes: 8 additions & 11 deletions Tests/test_decompression_bomb.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,16 @@


class TestDecompressionBomb:
def teardown_method(self) -> None:
Image.MAX_IMAGE_PIXELS = ORIGINAL_LIMIT

def test_no_warning_small_file(self) -> None:
# Implicit assert: no warning.
# A warning would cause a failure.
with Image.open(TEST_FILE):
pass

def test_no_warning_no_limit(self) -> None:
def test_no_warning_no_limit(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Arrange
# Turn limit off
Image.MAX_IMAGE_PIXELS = None
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", None)
assert Image.MAX_IMAGE_PIXELS is None

# Act / Assert
Expand All @@ -33,18 +30,18 @@ def test_no_warning_no_limit(self) -> None:
with Image.open(TEST_FILE):
pass

def test_warning(self) -> None:
def test_warning(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Set limit to trigger warning on the test file
Image.MAX_IMAGE_PIXELS = 128 * 128 - 1
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 128 * 128 - 1)
assert Image.MAX_IMAGE_PIXELS == 128 * 128 - 1

with pytest.warns(Image.DecompressionBombWarning):
with Image.open(TEST_FILE):
pass

def test_exception(self) -> None:
def test_exception(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Set limit to trigger exception on the test file
Image.MAX_IMAGE_PIXELS = 64 * 128 - 1
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 64 * 128 - 1)
assert Image.MAX_IMAGE_PIXELS == 64 * 128 - 1

with pytest.raises(Image.DecompressionBombError):
Expand All @@ -66,9 +63,9 @@ def test_exception_gif_extents(self) -> None:
with pytest.raises(Image.DecompressionBombError):
im.seek(1)

def test_exception_gif_zero_width(self) -> None:
def test_exception_gif_zero_width(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Set limit to trigger exception on the test file
Image.MAX_IMAGE_PIXELS = 4 * 64 * 128
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 4 * 64 * 128)
assert Image.MAX_IMAGE_PIXELS == 4 * 64 * 128

with pytest.raises(Image.DecompressionBombError):
Expand Down
29 changes: 13 additions & 16 deletions Tests/test_file_fli.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,22 +35,19 @@ def test_sanity() -> None:
assert im.is_animated


def test_prefix_chunk() -> None:
ImageFile.LOAD_TRUNCATED_IMAGES = True
try:
with Image.open(animated_test_file_with_prefix_chunk) as im:
assert im.mode == "P"
assert im.size == (320, 200)
assert im.format == "FLI"
assert im.info["duration"] == 171
assert im.is_animated

palette = im.getpalette()
assert palette[3:6] == [255, 255, 255]
assert palette[381:384] == [204, 204, 12]
assert palette[765:] == [252, 0, 0]
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
def test_prefix_chunk(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with Image.open(animated_test_file_with_prefix_chunk) as im:
assert im.mode == "P"
assert im.size == (320, 200)
assert im.format == "FLI"
assert im.info["duration"] == 171
assert im.is_animated

palette = im.getpalette()
assert palette[3:6] == [255, 255, 255]
assert palette[381:384] == [204, 204, 12]
assert palette[765:] == [252, 0, 0]


@pytest.mark.skipif(is_pypy(), reason="Requires CPython")
Expand Down
25 changes: 11 additions & 14 deletions Tests/test_file_ico.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,26 +243,23 @@ def test_draw_reloaded(tmp_path: Path) -> None:
assert_image_equal_tofile(im, "Tests/images/hopper_draw.ico")


def test_truncated_mask() -> None:
def test_truncated_mask(monkeypatch: pytest.MonkeyPatch) -> None:
# 1 bpp
with open("Tests/images/hopper_mask.ico", "rb") as fp:
data = fp.read()

ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
data = data[:-3]

try:
with Image.open(io.BytesIO(data)) as im:
assert im.mode == "1"
with Image.open(io.BytesIO(data)) as im:
assert im.mode == "1"

# 32 bpp
output = io.BytesIO()
expected = hopper("RGBA")
expected.save(output, "ico", bitmap_format="bmp")
# 32 bpp
output = io.BytesIO()
expected = hopper("RGBA")
expected.save(output, "ico", bitmap_format="bmp")

data = output.getvalue()[:-1]
data = output.getvalue()[:-1]

with Image.open(io.BytesIO(data)) as im:
assert im.mode == "RGB"
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
with Image.open(io.BytesIO(data)) as im:
assert im.mode == "RGB"
12 changes: 6 additions & 6 deletions Tests/test_file_jpeg.py
Original file line number Diff line number Diff line change
Expand Up @@ -530,12 +530,13 @@ def test_ff00_jpeg_header(self) -> None:
@mark_if_feature_version(
pytest.mark.valgrind_known_error, "libjpeg_turbo", "2.0", reason="Known Failing"
)
def test_truncated_jpeg_should_read_all_the_data(self) -> None:
def test_truncated_jpeg_should_read_all_the_data(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
filename = "Tests/images/truncated_jpeg.jpg"
ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with Image.open(filename) as im:
im.load()
ImageFile.LOAD_TRUNCATED_IMAGES = False
assert im.getbbox() is not None

def test_truncated_jpeg_throws_oserror(self) -> None:
Expand Down Expand Up @@ -1024,7 +1025,7 @@ def test_save_xmp(self, tmp_path: Path) -> None:
im.save(f, xmp=b"1" * 65505)

@pytest.mark.timeout(timeout=1)
def test_eof(self) -> None:
def test_eof(self, monkeypatch: pytest.MonkeyPatch) -> None:
# Even though this decoder never says that it is finished
# the image should still end when there is no new data
class InfiniteMockPyDecoder(ImageFile.PyDecoder):
Expand All @@ -1039,9 +1040,8 @@ def decode(
im.tile = [
ImageFile._Tile("INFINITE", (0, 0, 128, 128), 0, ("RGB", 0, 1)),
]
ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
im.load()
ImageFile.LOAD_TRUNCATED_IMAGES = False

def test_separate_tables(self) -> None:
im = hopper()
Expand Down
13 changes: 5 additions & 8 deletions Tests/test_file_jpeg2k.py
Original file line number Diff line number Diff line change
Expand Up @@ -181,14 +181,11 @@ def test_load_dpi() -> None:
assert "dpi" not in im.info


def test_restricted_icc_profile() -> None:
ImageFile.LOAD_TRUNCATED_IMAGES = True
try:
# JPEG2000 image with a restricted ICC profile and a known colorspace
with Image.open("Tests/images/balloon_eciRGBv2_aware.jp2") as im:
assert im.mode == "RGB"
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
def test_restricted_icc_profile(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
# JPEG2000 image with a restricted ICC profile and a known colorspace
with Image.open("Tests/images/balloon_eciRGBv2_aware.jp2") as im:
assert im.mode == "RGB"


@pytest.mark.skipif(
Expand Down
64 changes: 27 additions & 37 deletions Tests/test_file_png.py
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ def test_verify_struct_error(self) -> None:
with pytest.raises((OSError, SyntaxError)):
im.verify()

def test_verify_ignores_crc_error(self) -> None:
def test_verify_ignores_crc_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
# check ignores crc errors in ancillary chunks

chunk_data = chunk(b"tEXt", b"spam")
Expand All @@ -373,24 +373,20 @@ def test_verify_ignores_crc_error(self) -> None:
with pytest.raises(SyntaxError):
PngImagePlugin.PngImageFile(BytesIO(image_data))

ImageFile.LOAD_TRUNCATED_IMAGES = True
try:
im = load(image_data)
assert im is not None
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
im = load(image_data)
assert im is not None

def test_verify_not_ignores_crc_error_in_required_chunk(self) -> None:
def test_verify_not_ignores_crc_error_in_required_chunk(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
# check does not ignore crc errors in required chunks

image_data = MAGIC + IHDR[:-1] + b"q" + TAIL

ImageFile.LOAD_TRUNCATED_IMAGES = True
try:
with pytest.raises(SyntaxError):
PngImagePlugin.PngImageFile(BytesIO(image_data))
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with pytest.raises(SyntaxError):
PngImagePlugin.PngImageFile(BytesIO(image_data))

def test_roundtrip_dpi(self) -> None:
# Check dpi roundtripping
Expand Down Expand Up @@ -600,7 +596,7 @@ def test_roundtrip_private_chunk(self) -> None:
(b"prIV", b"VALUE3", True),
]

def test_textual_chunks_after_idat(self) -> None:
def test_textual_chunks_after_idat(self, monkeypatch: pytest.MonkeyPatch) -> None:
with Image.open("Tests/images/hopper.png") as im:
assert "comment" in im.text
for k, v in {
Expand All @@ -614,18 +610,17 @@ def test_textual_chunks_after_idat(self) -> None:
with pytest.raises(OSError):
assert isinstance(im.text, dict)

# Raises an EOFError in load_end
with Image.open("Tests/images/hopper_idat_after_image_end.png") as im:
assert im.text == {"TXT": "VALUE", "ZIP": "VALUE"}

# Raises a UnicodeDecodeError in load_end
with Image.open("Tests/images/truncated_image.png") as im:
# The file is truncated
with pytest.raises(OSError):
im.text
ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
assert isinstance(im.text, dict)
ImageFile.LOAD_TRUNCATED_IMAGES = False

# Raises an EOFError in load_end
with Image.open("Tests/images/hopper_idat_after_image_end.png") as im:
assert im.text == {"TXT": "VALUE", "ZIP": "VALUE"}

def test_unknown_compression_method(self) -> None:
with pytest.raises(SyntaxError, match="Unknown compression method"):
Expand All @@ -651,15 +646,16 @@ def test_padded_idat(self) -> None:
@pytest.mark.parametrize(
"cid", (b"IHDR", b"sRGB", b"pHYs", b"acTL", b"fcTL", b"fdAT")
)
def test_truncated_chunks(self, cid: bytes) -> None:
def test_truncated_chunks(
self, cid: bytes, monkeypatch: pytest.MonkeyPatch
) -> None:
fp = BytesIO()
with PngImagePlugin.PngStream(fp) as png:
with pytest.raises(ValueError):
png.call(cid, 0, 0)

ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
png.call(cid, 0, 0)
ImageFile.LOAD_TRUNCATED_IMAGES = False

@pytest.mark.parametrize("save_all", (True, False))
def test_specify_bits(self, save_all: bool, tmp_path: Path) -> None:
Expand Down Expand Up @@ -789,17 +785,14 @@ class MyStdOut:
with Image.open(mystdout) as reloaded:
assert_image_equal_tofile(reloaded, TEST_PNG_FILE)

def test_truncated_end_chunk(self) -> None:
def test_truncated_end_chunk(self, monkeypatch: pytest.MonkeyPatch) -> None:
with Image.open("Tests/images/truncated_end_chunk.png") as im:
with pytest.raises(OSError):
im.load()

ImageFile.LOAD_TRUNCATED_IMAGES = True
try:
with Image.open("Tests/images/truncated_end_chunk.png") as im:
assert_image_equal_tofile(im, "Tests/images/hopper.png")
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with Image.open("Tests/images/truncated_end_chunk.png") as im:
assert_image_equal_tofile(im, "Tests/images/hopper.png")


@pytest.mark.skipif(is_win32(), reason="Requires Unix or macOS")
Expand All @@ -808,19 +801,16 @@ class TestTruncatedPngPLeaks(PillowLeakTestCase):
mem_limit = 2 * 1024 # max increase in K
iterations = 100 # Leak is 56k/iteration, this will leak 5.6megs

def test_leak_load(self) -> None:
def test_leak_load(self, monkeypatch: pytest.MonkeyPatch) -> None:
with open("Tests/images/hopper.png", "rb") as f:
DATA = BytesIO(f.read(16 * 1024))

ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
with Image.open(DATA) as im:
im.load()

def core() -> None:
with Image.open(DATA) as im:
im.load()

try:
self._test_leak(core)
finally:
ImageFile.LOAD_TRUNCATED_IMAGES = False
self._test_leak(core)
5 changes: 2 additions & 3 deletions Tests/test_file_tiff.py
Original file line number Diff line number Diff line change
Expand Up @@ -939,11 +939,10 @@ def test_string_dimension(self) -> None:

@pytest.mark.timeout(6)
@pytest.mark.filterwarnings("ignore:Truncated File Read")
def test_timeout(self) -> None:
def test_timeout(self, monkeypatch: pytest.MonkeyPatch) -> None:
with Image.open("Tests/images/timeout-6646305047838720") as im:
ImageFile.LOAD_TRUNCATED_IMAGES = True
monkeypatch.setattr(ImageFile, "LOAD_TRUNCATED_IMAGES", True)
im.load()
ImageFile.LOAD_TRUNCATED_IMAGES = False

@pytest.mark.parametrize(
"test_file",
Expand Down
Loading

0 comments on commit bf94539

Please sign in to comment.