import png
from exploits.bombs.images import png as png_exploit
from test.exploits.dummy_output import DummyOutput


def test_small_png_bomb():
    output = DummyOutput()
    target_width = 20
    target_height = 10
    png_exploit.options['width'] = target_width
    png_exploit.options['height'] = target_height
    png_exploit.run(output)
    assert len(output) == 1
    width, height, _, _ = png.Reader(bytes=output[0]).read()
    assert width == target_width
    assert height == target_height


def test_large_png_bomb():
    output = DummyOutput()
    target_width = 200000
    target_height = 100000
    png_exploit.options['width'] = target_width
    png_exploit.options['height'] = target_height
    png_exploit.run(output)
    assert len(output) == 1
    width, height, _, _ = png.Reader(bytes=output[0]).read()
    assert width == target_width
    assert height == target_height
