client/pm3_eml_mfd_test.py

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