from unittest import TestCase

from Ciphers.types import Action
from Ciphers.vigenere import use_autokey, use_vigenere


class VigenereTestCase(TestCase):
    def test_vigenere_encryption(self):
        self.assertEqual('DLGO WJ D WIQOOXH', use_vigenere('This is a message!', 'keyword', action=Action.ENCRYPT))
        self.assertEqual('DLGO WJ D WIQOOXH!', use_vigenere('This is a message!', 'keyword',
                                                            chars_to_keep=' !', action=Action.ENCRYPT))
        self.assertEqual('DLGOWJDWIQOOXH', use_vigenere('This is a message!', 'keyword',
                                                        chars_to_keep='', action=Action.ENCRYPT))
        self.assertEqual('MOQK SW Y ISJVIYI', use_vigenere('This is a message!', 'This keyword is even longer than msg',
                                                           action=Action.ENCRYPT))

    def test_vigenere_decryption(self):
        self.assertEqual('THIS IS A MESSAGE', use_vigenere('Dlgo wj d wiqooxh!', 'keyword', action=Action.DECRYPT))
        self.assertEqual('THIS IS A MESSAGE!', use_vigenere('Dlgo wj d wiqooxh!', 'keyword',
                                                            chars_to_keep=' !', action=Action.DECRYPT))
        self.assertEqual('THISISAMESSAGE', use_vigenere('Dlgo wj d wiqooxh!', 'keyword',
                                                        chars_to_keep='', action=Action.DECRYPT))
        self.assertEqual('THIS IS A MESSAGE', use_vigenere('Moqk sw y isjviyi', 'This keyword is even longer than msg',
                                                           action=Action.DECRYPT))
        self.assertEqual('THISISAMESSAGE', use_vigenere('Moqk sw y isjviyi', 'This keyword is even longer than msg',
                                                        chars_to_keep='', action=Action.DECRYPT))

    def test_autokey_encryption(self):
        self.assertEqual('SLJJ IL H AWKJEL QGJWTSI',
                         use_autokey('This is a secret message', 'zebra', action=Action.ENCRYPT))

    def test_autokey_decryption(self):
        self.assertEqual('THIS IS A SECRET MESSAGE',
                         use_autokey('SLJJ IL H AWKJEL QGJWTSI', 'zebra'))