TST: add test of truncated files and buffers

tacaswell · tacaswell · commit a17787b93e76 · 2017-10-01T20:52:03.000-07:00
diff --git a/lib/matplotlib/tests/test_png.py b/lib/matplotlib/tests/test_png.py
@@ -2,16 +2,17 @@
                         unicode_literals)
 
 import six
-
+from six import BytesIO
 import glob
 import os
-
+import shutil
 import numpy as np
+import pytest
 
 from matplotlib.testing.decorators import image_comparison
 from matplotlib import pyplot as plt
 import matplotlib.cm as cm
-
+import tempfile
 import sys
 on_win = (sys.platform == 'win32')
 
@@ -46,3 +47,30 @@ def test_imread_png_uint16():
 
     assert (img.dtype == np.uint16)
     assert np.sum(img.flatten()) == 134184960
+
+
+def test_truncated_file():
+    d = tempfile.mkdtemp()
+    fname = os.path.join(d, 'test.png')
+    fname_t = os.path.join(d, 'test_truncated.png')
+    plt.savefig(fname)
+    with open(fname, 'rb') as fin:
+        buf = fin.read()
+    with open(fname_t, 'wb') as fout:
+        fout.write(buf[:20])
+
+    with pytest.raises(Exception):
+        plt.imread(fname_t)
+
+    shutil.rmtree(d)
+
+
+def test_truncated_buffer():
+    b = BytesIO()
+    plt.savefig(b)
+    b.seek(0)
+    b2 = BytesIO(b.read(20))
+    b2.seek(0)
+
+    with pytest.raises(Exception):
+        plt.imread(b2)
