test_encrypt.py

import unittest
import time
import os
import tempfile
import shutil


class MainTest(unittest.TestCase):
    def setUp(self):
        self.tic = time.time()  # todo this is sooo ugly
        self.test_dir = tempfile.mkdtemp()

    def tearDown(self):
        self.toc = time.time()
        t = self.toc - self.tic
        print('Time: %.3f' % (t))

        # Clean up
        shutil.rmtree(self.test_dir)

    def test_encrypt_scan_single(self):
        import encrypt

        expected_files = [
            '123001_LastnameA_aes.pdf']

        # Prepare parameter
        in_pdf = './tests/assets/pdfs/123001_LastnameA.pdf'
        enc_pdf = '123001_LastnameA_aes.pdf'

        out_dir = os.path.join(self.test_dir, 'out')
        os.mkdir(out_dir)

        # Encrypt files
        encrypt.encrypt(
            pdf_file=in_pdf, enc_file=os.path.join(out_dir, enc_pdf),
            password='tests_are_fun')

        created_files = os.listdir(out_dir)
        created_files.sort()
        self.assertEqual(expected_files, created_files)

    def test_encrypt_scans(self):
        import encrypt

        expected_files = ["123001_LastnameA_aes.pdf",
                          "123002_LastnameB_aes.pdf",
                          "123010_LastnameJ_aes.pdf",
                          "123011_LastnameK_aes.pdf",
                          'passwords.csv']

        # Prepare parameter
        in_dir = './tests/assets/pdfs'

        out_dir = os.path.join(self.test_dir, 'out')
        os.mkdir(out_dir)

        # Encrypt files
        encrypt.main([in_dir, out_dir])

        created_files = os.listdir(out_dir)
        created_files.sort()
        self.assertEqual(expected_files, created_files)

    def test_encrypt_supplements(self):
        import encrypt

        expected_files = ['GDET3_20H_aes.pdf', 'GDET3_20H_loes_aes.pdf',
                          'passwords.csv']

        # Prepare parameter
        in_dir = './tests/assets/supplements'

        out_dir = os.path.join(self.test_dir, 'out')
        os.mkdir(out_dir)

        # Encrypt files
        encrypt.main([in_dir, out_dir])

        created_files = os.listdir(out_dir)
        created_files.sort()
        self.assertEqual(expected_files, created_files)