From 7a1ca875f0a2cf9eac9f9824dc0b8c807d07b63b Mon Sep 17 00:00:00 2001
From: uchood <ucagoog@gmail.com>
Date: Mon, 29 Jan 2018 02:16:16 +0300
Subject: [PATCH] Add tests iterable: sum streams exceed zip64limit with not
 allowed zip64 add tests iterable: sum streams exceed zip64limit with allowed
 zip64

---
 tests/test_zipstream.py | 99 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/tests/test_zipstream.py b/tests/test_zipstream.py
index 9910fe2..2c891a3 100644
--- a/tests/test_zipstream.py
+++ b/tests/test_zipstream.py
@@ -96,5 +96,104 @@ def test_write_iterable_no_archive(self):
         z = zipstream.ZipFile(mode='w')
         self.assertRaises(TypeError, z.write_iter, iterable=range(10))
 
+    def test_write_iterable_zip64_with_not_allow_zip64_many_smalls(self):
+        # check many small streams that sum length require ZIP64 extensions when not allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=False)
+
+        def string_small_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter + len_sample < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_small_generator(), string_small_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            self.assertRaises(zipfile.LargeZipFile, lambda: [f.write(c) for c in z])
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_not_allow_zip64_1_big_file(self):
+        # check 1 big stream that length require ZIP64 extensions when not allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=False)
+
+        def string_big_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_big_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            self.assertRaises(zipfile.LargeZipFile, lambda: [f.write(c) for c in z])
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_allow_zip64_many_smalls(self):
+        # check many small streams that sum length require ZIP64 extensions when allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=True)
+
+        def string_small_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter + len_sample < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_small_generator(), string_small_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            for chunk in z:
+                f.write(chunk)
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+    def test_write_iterable_zip64_with_allow_zip64_1_big_file(self):
+        # check 1 big stream that length require ZIP64 extensions when allowed zip64
+        z = zipstream.ZipFile(mode='w', allowZip64=True)
+
+        def string_big_generator():
+            counter = 0
+            sample = b'zipstream0' * 10000000
+            len_sample = len(sample)
+            while counter < zipstream.ZIP64_LIMIT:
+                counter += len_sample
+                yield sample
+
+        data = [string_big_generator()]
+        for i, d in enumerate(data):
+            z.write_iter(iterable=d, arcname='data_{0}'.format(i))
+        f = tempfile.NamedTemporaryFile(suffix='zip', delete=False)
+        try:
+            for chunk in z:
+                f.write(chunk)
+            f.close()
+        except Exception:
+            raise
+        finally:
+            os.remove(f.name)
+
+
 if __name__ == '__main__':
     unittest.main()