From 4dafc2696a39b3b892cd4c017326229c587c10ff Mon Sep 17 00:00:00 2001
From: Vadim Lebedev <vadim@mbdsys.com>
Date: Thu, 11 Aug 2022 00:37:32 +0200
Subject: [PATCH] Implement testing of password protected zip creation

---
 tests.py | 59 +++++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 52 insertions(+), 7 deletions(-)

diff --git a/tests.py b/tests.py
index 6c6ad7c..9380e63 100644
--- a/tests.py
+++ b/tests.py
@@ -266,22 +266,58 @@ ZIP_CONTENT='UEsDBAoACQAAAKubCVVjZ7b1FQAAAAkAAAAIABwAdGVzdC50eHRVVAkAA5K18mKStfJ
 ITEM_CONTENT='test.txt\n'
 ITEM_NAME='test.txt'
 
+ZIP1_PWD='pwd'
+ZIP2_PWD='12345'
+def create_file_from_content():
+    if PY3:
+        with open(ZIPPATH, mode='wb') as f:
+            f.write(base64.b64decode(ZIP_CONTENT))
+    else:
+        with open(ZIPPATH, mode='w') as f:
+            f.write(base64.b64decode(ZIP_CONTENT))
 
-class TestPasswordProtection(unittest.TestCase):
+
+def create_protected_zip():
+    z = ZipFile(ZIPPATH, mode='w', password=ZIP2_PWD)
+    z.writestr(ITEM_NAME, ITEM_CONTENT)
+    z.close()
+
+
+class TestProtectedReading(unittest.TestCase):
     def setUp(self):
+        create_file_from_content()
+
+
+    def tearDown(self):
+        os.remove(ZIPPATH)
+
+    def test_read_with_password(self):
+        z = ZipFile(ZIPPATH, 'r', password=ZIP1_PWD)
         if PY3:
-            with open(ZIPPATH, mode='wb') as f:
-                f.write(base64.b64decode(ZIP_CONTENT))
+            self.assertEqual(z.read(ITEM_NAME), bytes(ITEM_CONTENT, 'utf-8'))
         else:
-            with open(ZIPPATH, mode='w') as f:
-                f.write(base64.b64decode(ZIP_CONTENT))
+            self.assertEqual(z.read(ITEM_NAME), ITEM_CONTENT)
+        z.close()
+
+    def test_read_without_password(self):
+        z = ZipFile(ZIPPATH, 'r')
+        self.assertRaises(RuntimeError, z.read, ITEM_NAME)
+        z.close()
+
+    def test_read_with_wrong_password(self):
+        z = ZipFile(ZIPPATH, 'r', password='wrong')
+        self.assertRaises(RuntimeError, z.read, ITEM_NAME)
+        z.close()
+
+class TestProtectedWriting(unittest.TestCase):
+    def setUp(self):
+        create_protected_zip()
 
-    
     def tearDown(self):
         os.remove(ZIPPATH)
 
     def test_read_with_password(self):
-        z = ZipFile(ZIPPATH, 'r', password='pwd')
+        z = ZipFile(ZIPPATH, 'r', password=ZIP2_PWD)
         if PY3:
             self.assertEqual(z.read(ITEM_NAME), bytes(ITEM_CONTENT, 'utf-8'))
         else:
@@ -298,6 +334,15 @@ class TestPasswordProtection(unittest.TestCase):
         self.assertRaises(RuntimeError, z.read, ITEM_NAME)
         z.close()
 
+    def test_read_with_password_list(self):
+        z = ZipFile(ZIPPATH, 'r', password=[ZIP1_PWD, ZIP2_PWD])
+        if PY3:
+            self.assertEqual(z.read(ITEM_NAME), bytes(ITEM_CONTENT, 'utf-8'))
+        else:
+            self.assertEqual(z.read(ITEM_NAME), ITEM_CONTENT)
+        z.close()
+
+
 
 class TestHighLevelAPI(unittest.TestCase):
     def setUp(self):