1 files changed, 50 insertions, 0 deletions
diff --git a/tests/test_crypto.py b/tests/test_crypto.py
new file mode 100644
index 0000000..83fbe41
--- /dev/null
+++ b/tests/test_crypto.py
@@ -0,0 +1,50 @@
+# test_crypto.py
+import unittest
+from io import BytesIO
+from cryptography.fernet import Fernet
+from crypto import generate_key, load_key, encrypt, decrypt
+
+class TestCryptoFunctions(unittest.TestCase):
+
+    def setUp(self):
+        self.key = generate_key()
+        with open('vault.sqlite', 'wb') as vault:
+            vault.write(b'Hello, world!')
+
+    def tearDown(self):
+        try:
+            os.remove('vault.sqlite')
+        except FileNotFoundError:
+            pass
+
+    def test_generate_key(self):
+        new_key = generate_key()
+        self.assertNotEqual(new_key, self.key)
+
+    def test_load_key(self):
+        key_file = 'key.bin'
+        with open(key_file, 'wb') as key:
+            key.write(self.key)
+        loaded_key = load_key(key_file)
+        self.assertEqual(loaded_key, self.key)
+
+    def test_encrypt_decrypt(self):
+        encrypt(self.key, filename='vault.sqlite')
+        with open('vault.sqlite', 'rb') as vault:
+            encrypted_data = vault.read()
+        decrypt(self.key)
+        with open('vault.sqlite', 'rb') as vault:
+            decrypted_data = vault.read()
+        f = Fernet(self.key)
+        self.assertEqual(f.decrypt(encrypted_data), b'Hello, world!')
+
+    def test_load_invalid_key(self):
+        key_file = 'key.bin'
+        try:
+            load_key(key_file)
+            self.fail('Expected FileNotFoundError')
+        except FileNotFoundError:
+            pass
+
+if __name__ == '__main__':
+    unittest.main()