<filename>test_main.py
#!/usr/bin/env python3
'''Author: c0llision '''
import unittest
from main import *


class Test(unittest.TestCase):
    def test(self):
        password = '<PASSWORD>'
        secret = 'secret message'

        ciphertext = encrypt_data(secret, password)
        cleartext = decrypt_data(ciphertext, password)

        self.assertEqual(secret, cleartext)

    def test_scrypt(self):
        password = '<PASSWORD>'
        salt = b'123456'
        kdf_algo = 'scrypt'
        iterations = 14
        expected_key = b'\<KEY>'

        key = derive_key(password, salt, kdf_algo, iterations)
        self.assertEqual(key, expected_key)

    def test_pbkdf(self):
        password = '<PASSWORD>'
        salt = b'123456'
        kdf_algo = 'pbkdf'
        iterations = 10000
        expected_key = b'\<KEY>'

        key = derive_key(password, salt, kdf_algo, iterations)
        self.assertEqual(key, expected_key)


if __name__ == '__main__':
    unittest.main()