python 單元測試unittest

最近因爲項目學習並運用了python的單元測試框架，豐富知識儲備。直接上代碼，主要是做一個加解密模塊的測試，直接上代碼

# -*- coding: utf-8 -*-
import os
import sys
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto import Random
import unittest

class kdba_encryt(object):
    """
    A class grouping together all of the Crypto API functions.
    hers: 'AES'
    Modes: 'ECB', 'CBC', 'CFB', 'OFB', 'CTR'
    Hash Functions: 'SHA256'
    """
    def __init__(self, 
                 chunksize = 64*1024,
                 type = 'AES', 
                 mode = 'CBC', 
                 password = 'kdba_debug',
                 IsRemove=False,
                 IsCover=False):
        self.type = type
        self.mode = mode
        self.chunksize = chunksize
        self.IsRemove = IsRemove
        self.IsCover = IsCover
        self.key = self._getkey(password)
		  # random produce distinct cipher text (initialization vector)
        self.iv = Random.new().read(16)
	
    	  # default enable AES and CBC mode, if want use other algorthim 
		  # define it in this place	
        if self.type == 'AES' and self.mode == 'CBC':
            self.encryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
            self.decryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
        else:
            pass
		
        self._print_describe()

    def _print_describe(self):
        print("Create encrypto and decrypto operator")
        print("Ciphers:{}\nMode:{}\n".format(self.type, self.mode))
		  #print("Key:%#x" % self.key.encode())
		  #print("IV:%#x" % self.iv)
		

    def _getkey(self, password):
        hasher= SHA256.new(password.encode('utf-8'))
        return hasher.digest()	

    def encrypt(self, input_file, output_file):
        
#	    zfill returns the string with 0 filled at start
        filesize = str(os.path.getsize(input_file)).zfill(16)
	
#	     encryptor = AES.new(self.key,AES.MODE_CBC, self.iv)
        with open(input_file,'rb') as infile:
            with open(output_file,'wb') as outfile:
                outfile.write(filesize.encode('utf-8'))
                while True:
                    chunk = infile.read(self.chunksize)
	
                    if len(chunk) == 0:
                        break
                    elif len(chunk) % 16 != 0:
                        # Padding 0 to align 16 byte
                        chunk += b' '  * (16 - (len(chunk) % 16))
	
                    outfile.write(self.encryptor.encrypt(chunk))
        
        if self.IsRemove:
            if os.path.exists(input_file):
                os.unlink(input_file)

    def decrypt(self, src_file, dst_file):
        with open(src_file,'rb') as infile:
            filesize = int(infile.read(16))  # int type is safe
	
            #decryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
	
            with open(dst_file,'wb') as outfile:
                while True:
                    chunk = infile.read(self.chunksize)
                    if len(chunk) == 0:
                        break
                    outfile.write(self.decryptor.decrypt(chunk))
                outfile.truncate(filesize)

class TestEncryptDecrypt(unittest.TestCase):
    """
    unit test foro encrypt and decrypt
    """
    def setUp(self):
        """initial some file name"""
        print('initial...')
        self.filename = './model0.txt'
        self.old_filename = './model0_old.txt'
        
    def _read_file(self):
        print('load file...')
        """load encrypt file into byte"""
        with open(self.filename,'rb') as fh:
            file_byte = fh.read()
        return file_byte
    
    def _read_old_file(self):
        print('load decrypt file')
        """load decrypt file into byte"""
        with open(self.old_filename,'rb') as fo:
            old_file_byte = fo.read()
        return old_file_byte
    
    def test_equal(self):
        file_byte = self._read_file()
        old_file_byte = self._read_old_file()
        self.assertEqual(file_byte,old_file_byte)
        
    @unittest.skip("This method don't use now!")
    def test_size(self):
        file_byte = self._read_file()
        self.assertLessEqual(2**20,sys.getsizeof(file_byte))
      
    def remove_file(self):
        pass
    
    def tearDown(self):
        pass
        

if __name__== '__main__':
#    kdba_encryt = kdba_encryt()
#    kdba_encryt.encrypt('./SecureFile.py', './SecureFile_enc.py')
#    kdba_encryt.decrypt('./SecureFile_enc.py', './SecureFile_old.py')
#    kdba_encryt.encrypt('./test.py', './test_enc.py')
#    kdba_encryt.decrypt('./test_enc.py', './test_old.py')
#    kdba_encryt.encrypt('./crp.txt','./crp_enc.txt')
#    kdba_encryt.decrypt('./crp_enc.txt','./crp_old.txt')
    
    # decide the order of execution
#    unittest.main()
    suite = unittest.TestSuite()
    tests = [TestEncryptDecrypt('test_equal'),TestEncryptDecrypt('test_size')]
    suite.addTests(tests)
    with open('UnittestTextReport.txt','a+') as f:    
#        runner = unittest.TextTestRunner(stream=f,verbosity=1)
        runner = unittest.TextTestRunner()
        runner.run(suite)

參考資料：https://huilansame.github.io/huilansame.github.io/archivers/python-unittest

https://github.com/coleifer/beefish/blob/master/beefish.py

https://docs.python.org/zh-cn/3/library/unittest.html#assert-methods