added tests for encryptor.py

gradio/encryptor.py

@@ -19,6 +19,7 @@ def encrypt(
     IV = Random.new().read(AES.block_size)  # generate IV
     encryptor = AES.new(key, AES.MODE_CBC, IV)
     padding = AES.block_size - len(source) % AES.block_size  # calculate needed padding
+    print(type(source), type(padding))
     source += bytes([padding]) * padding  # Python 2.x: source += chr(padding) * padding
     data = IV + encryptor.encrypt(source)  # store the IV at the beginning and encrypt
     return data

test/test_encryptor.py

@@ -0,0 +1,33 @@
+import os
+import unittest
+
+from gradio import encryptor, processing_utils
+from gradio.test_data import BASE64_IMAGE
+
+
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+
+
+class TestKeyGenerator(unittest.TestCase):
+    def test_same_pass(self):
+        key1 = encryptor.get_key("test")
+        key2 = encryptor.get_key("test")
+        self.assertEquals(key1, key2)        
+
+    def test_diff_pass(self):
+        key1 = encryptor.get_key("test")
+        key2 = encryptor.get_key("diff_test")
+        self.assertNotEquals(key1, key2)        
+
+
+class TestEncryptorDecryptor(unittest.TestCase):
+    def test_same_pass(self):
+        key = encryptor.get_key("test")
+        data, _ = processing_utils.decode_base64_to_binary(BASE64_IMAGE)
+        encrypted_data = encryptor.encrypt(key, data)
+        decrypted_data = encryptor.decrypt(key, encrypted_data)
+        self.assertEquals(data, decrypted_data)
+        
+
+if __name__ == "__main__":
+    unittest.main()