from unittest import TestCase
from benchmark.railfence_cipher import railencrypt
from benchmark.railfence_cipher import raildecrypt


class Test_railencrypt(TestCase):
    #  distances_true = {1: [0, 0, 0, 0, 0, 1, 1, 1, 1], 2: [4, 3, 2, 1, 0], 3: [3, 2, 1, 0], 4: [0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1]}
    #  distances_false = {1: [1, 1, 1, 1, 1, 0, 0, 0, 0], 2: [0, 0, 0, 0, 1], 3: [0, 0, 0, 1], 4: [80, 0, 0, 0, 0, 0, 0, 0, 65, 0, 113, 0, 0, 0, 0, 0, 69, 0, 0, 0, 109, 0, 0, 0, 111, 0, 0, 0, 0, 0, 73, 0, 98, 0, 0, 0, 0, 0, 0, 0, 37, 0, 0, 0, 0]}
    def test_railencrypt_1(self):
        assert railencrypt(st='PqmI%boEA', k=5) == 'PAqEmoIb%'


class Test_raildecrypt(TestCase):
    #  distances_true = {5: [0, 0, 0, 1, 1, 0, 0, 1], 6: [2, 1, 0, 1, 0], 7: [1, 0, 1], 8: [0, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 1], 9: [0, 0, 0, 0, 0, 0, 0, 0], 10: [8, 0, 0, 1, 1, 0, 0, 1], 12: [0, 1, 0, 1], 11: [1, 0, 1, 0]}
    #  distances_false = {5: [1, 1, 1, 0, 0, 1, 1, 0], 6: [0, 0, 1, 0, 1], 7: [0, 1, 0], 8: [1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0], 9: [114, 32, 42, 113, 94, 101, 66, 91], 10: [0, 1, 1, 0, 0, 1, 1, 0], 12: [1, 0, 1, 0], 11: [0, 1, 0, 1]}
    def test_raildecrypt_1(self):
        assert raildecrypt(st='r^ qe[*B', k=3) == 'r *q^eB['