client/pyscripts/pm3_eml_mfd_test.py

   1 #!/usr/bin/env python3
   2
   3 from tempfile import mkdtemp
   4 from shutil import rmtree
   5 from string import hexdigits
   6 import unittest, os
   7 import pm3_eml2mfd, pm3_mfd2eml
   8
   9 class TestEmlMfd(unittest.TestCase):
  10     def setUp(self):
  11         self.tmpdir = mkdtemp()
  12
  13     def tearDown(self):
  14         rmtree(self.tmpdir)
  15
  16     EML2MFD_TESTCASES = [
  17             ('', ''),
  18             ("41424344\r\n45464748\n494A4B4C\n", "ABCDEFGHIJKL")
  19             ]
  20     def test_eml2mfd(self):
  21         self.three_argument_test(pm3_eml2mfd.main, self.EML2MFD_TESTCASES)
  22
  23     def test_mfd2eml(self):
  24         self.three_argument_test(pm3_mfd2eml.main,
  25                 map(reversed, self.EML2MFD_TESTCASES), c14n=hex_c14n)
  26
  27     def three_argument_test(self, operation, cases, c14n=str):
  28         for case_input, case_output in cases:
  29             try:
  30                 inp_name = os.path.join(self.tmpdir, 'input')
  31                 out_name = os.path.join(self.tmpdir, 'output')
  32                 with open(inp_name, 'w') as in_file:
  33                     in_file.write(case_input)
  34                 operation(['', inp_name, out_name])
  35                 with open(out_name, 'r') as out_file:
  36                     self.assertEqual(c14n(case_output), c14n(out_file.read()))
  37             finally:
  38                 for file_name in inp_name, out_name:
  39                     if os.path.exists(file_name):
  40                         os.remove(file_name)
  41
  42
  43 def hex_c14n(inp):
  44     """
  45     Canonicalizes the input string by removing non-hexadecimal
  46     characters and making everything uppercase
  47     """
  48     return ''.join(c.upper() for c in inp if c in hexdigits)
  49
  50 if __name__ == '__main__':
  51     unittest.main()